from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np
import slrealizer.utils.utils as utils
import slrealizer.utils.constants as constants
import pandas as pd
import random
import galsim

class SLRealizer(object):

    """
    Worker class equipped with functions
    for making LSST-like "Object" and "Source" tables. These are
    inherited by child classes which are associated with a specific non-LSST catalog,
    e.g. the child class OM10Realizer realizes the lenses in the OM10 catalog into mock LSST tables.
    :class:`SLRealizer` objects are instantiated with
    an observation history, provided as a `pandas` dataframe.
    """

    def __init__(self, observation, add_moment_noise, add_flux_noise):
        # Observation history
        self.observation = observation
        self.num_obs = len(self.observation)

        # GalSim drawImage params
        self.fft_params = galsim.GSParams(maximum_fft_size=10240)
        self.pixel_scale = 0.1
        self.nx, self.ny = 49, 49

        # Source table df
        self.source_table = None
        # Source table column list
        self.source_columns = ['MJD', 'ccdVisitId', 'objectId', 'filter', 'psf_fwhm', 'x', 'y', 'apFlux', 'apFluxErr', 'apMag', 'apMagErr', 'trace', 'e1', 'e2', 'e_final', 'phi_final', ]

        # Controlling randomness
        self.add_moment_noise = add_moment_noise
        self.add_flux_noise = add_flux_noise
        if not self.add_flux_noise and not self.add_moment_noise:
            self.remove_random = True
        else:
            self.remove_random = False
        self.seed = 123
        np.random.seed(self.seed)

    def get_obs_info(self, obsID=None, rownum=None):
        if obsID is not None and rownum is not None:
            raise ValueError("Need to define either obsID or rownum, not both.")
        if obsID is not None:
            return self.observation.loc[self.observation['obsHistID']==obsID]
        elif rownum is not None:
            return self.observation.loc[rownum]

    def _get_lens_info(self, lensID):
        ''' This function will depend on the format of each lens catalog.
        '''
        raise NotImplementedError

    def draw_system(self, lens_info, obs_info, save_path=None):
        '''
        Draws all objects of the given lens system
        in the given observation conditions, using GalSim.

        Parameters
        ==========
        lens_info: pandas.Dataframe
            a row of the OM10 DB
        obs_info: pandas.Dataframe
            a row of the observation history df
        save_path: string
            path in which to save the image

        Returns
        =======
        galsim_img: galsim.Image
            A GalSim image object of the aggregate system
        '''
        # Old docstring, for comparison - to be deleted!
        # For more help writing nice docstrings, check out
        # https://python-sprints.github.io/pandas/guide/pandas_docstring.html
        '''
        Draws all objects of the given lens system
        in the given observation conditions, using GalSim.

        Keyword arguments:
        lens_info -- a row of the OM10 DB
        obs_info -- a row of the observation history df
        save_path -- path in which to save the image

        Returns:
        A GalSim object of the aggregate system used to render
        the image
        '''
        histID, MJD, band, PSF_FWHM, sky_mag = obs_info

        # Construct a "scene" of image components:
        # i) Lens galaxy #half_light_radius=lens_info['half_light_radius'],\
        scene = galsim.Gaussian(sigma=1.0, flux=lens_info['flux'])\
                       .shear(e=lens_info['e'], beta=lens_info['beta'])
        # ii) Lensed quasar images
        for i in xrange(lens_info['num_objects']):
            quasar = galsim.Gaussian(flux=lens_info['flux_'+str(i)], sigma=1.e-5)\
                         .shift(lens_info['xy_'+str(i)])
            scene += quasar

        # Convolve the scene with the PSF:
        psf = galsim.Gaussian(flux=1.0, fwhm=PSF_FWHM)
        galsim_obj = galsim.Convolve([scene, psf], gsparams=self.fft_params)
        galsim_img = galsim_obj.drawImage(nx=self.nx, ny=self.ny, scale=self.pixel_scale, method='no_pixel')
        if save_path is not None:
            plt.imshow(galsim_img.array, interpolation='none', aspect='auto')
            plt.savefig(save_path)
            plt.close()

        return galsim_img

    def estimate_parameters(self, galsim_img, method="raw_numerical"):
        """
        Performs shape estimati on on the galsim_img
        using either GalSim's HSM shape estimator or
        a native numerical moment calculator
        under the observation conditions obs_info

        Keyword arguments:
        galsim_img -- GalSim's Image object on which parameters will be estimated
        method -- one of "hsm" (GalSim's HSM shape estimator) or
                  "raw_numerical" (a native numerical moment calculator) [default: "raw"]

        Returns
        a dictionary of the lens properties,
        which can be used to draw the emulated image
        """
        estimated_params = {}
        if method == "hsm":
            try:
                shape_info = galsim_img.FindAdaptiveMom(guess_sig=self.pixel_scale*10.0)
            except:
                #print "HSM failed"
                return None

            # Calculate the real position from the arbitrary pixel position
            pixelCenter = galsim.PositionD(x=shape_info.moments_centroid.x, y=shape_info.moments_centroid.y)
            estimated_params['x'], estimated_params['y'] = utils.pixel_to_physical(shape_info.moments_centroid.x, self.nx, self.pixel_scale),\
                                             utils.pixel_to_physical(shape_info.moments_centroid.y, self.ny, self.pixel_scale)

            estimated_params['apFlux'] = float(np.sum(galsim_img.array))
            if self.DEBUG:
                estimated_params['hlr'] = galsim_img.calculateHLR(center=pixelCenter)
                estimated_params['det'] = (shape_info.moments_sigma * self.pixel_scale)**4.0
            estimated_params['trace'] = 2.0*galsim_img.calculateMomentRadius(center=pixelCenter, rtype='trace')**2.0
            estimated_params['e1'] = shape_info.observed_shape.e1
            estimated_params['e2'] = shape_info.observed_shape.e2
            estimated_params['e_final'] = shape_info.observed_shape.e
            estimated_params['phi_final'] = shape_info.observed_shape.beta
        elif method == "raw_numerical":
            image_array = galsim_img.array
            Ix, Iy = utils.get_first_moments_from_image(image_array, self.pixel_scale)
            Ixx, Ixy, Iyy = utils.get_second_moments_from_image(image_array, self.pixel_scale)
            estimated_params['apFlux'] = np.sum(image_array)
            estimated_params['x'] = Ix
            estimated_params['y'] = Iy
            trace = Ixx + Iyy
            estimated_params['e1'] = (Ixx - Iyy)/trace
            estimated_params['e2'] = 2.0*Ixy/trace
            estimated_params['trace'] = trace
            estimated_params['e_final'], estimated_params['phi_final'] = utils.e1e2_to_ephi(estimated_params['e1'], estimated_params['e2'])
        else:
            raise ValueError("Please enter a valid method, either 'hsm' or 'raw_numerical'")

        return estimated_params

    def draw_emulated_system(self, estimated_params):
        """
        Draws the emulated system, i.e. draws the aggregate system
        from properties HSM derived from the image, which was in turn
        drawn from the catalog's truth properties.
        Only runs when DEBUG == True.

        Returns
        a GalSim Image object of the emulated system
        """
        if not self.DEBUG:
            raise ValueError("Only runs in debug mode")
        system = galsim.Gaussian(flux=estimated_params['apFlux'], half_light_radius=estimated_params['hlr'])\
                       .shift(float(estimated_params['x']), float(estimated_params['y']))\
                       .shear(e1=estimated_params['e1'], e2=estimated_params['e2'])
        emulatedImg = system.drawImage(nx=self.nx, ny=self.ny, scale=self.pixel_scale, method='no_pixel')
        return emulatedImg

    def create_source_row(self, derived_params, objectId, obs_info):
        '''
        Returns a dictionary of lens system's properties
        computed the image of one lens system and the observation conditions,
        which makes up a row of the source table.

        Keyword arguments:
        derived_params -- derived lens properties
        obs_info -- a row of the observation history df

        Returns
        A dictionary with properties derived from HSM estimation
        (See code for which properties)
        '''
        histID, MJD, band, PSF_FWHM, sky_mag = obs_info

        derived_params['apFluxErr'] = utils.mag_to_flux(sky_mag-22.5)/5.0 # because Fb = 5 \sigma_b
        if self.add_moment_noise:
            derived_params['trace'] += utils.add_noise(mean=constants.get_second_moment_err(),
                                               stdev=constants.get_second_moment_err_std(),
                                               measurement=derived_params['trace'])
            derived_params['x'] += utils.add_noise(mean=constants.get_first_moment_err(),
                                           stdev=constants.get_first_moment_err_std(),
                                           measurement=derived_params['x'])
            derived_params['y'] += utils.add_noise(mean=constants.get_first_moment_err(),
                                           stdev=constants.get_first_moment_err_std(),
                                           measurement=derived_params['y'])
        if self.add_flux_noise:
            derived_params['apFlux'] += utils.add_noise(mean=0.0,
                                                stdev=derived_params['apFluxErr']) # flux rms not skyErr
        derived_params['apMag'] = utils.flux_to_mag(derived_params['apFlux'], from_unit='nMgy')
        derived_params['apMagErr'] = (2.5/np.log(10.0)) * derived_params['apFluxErr']/derived_params['apFlux']

        row = {'MJD': MJD, 'ccdVisitId': histID, 'filter': band, 'x': derived_params['x'], 'y': derived_params['y'],
               'apFlux': derived_params['apFlux'], 'apFluxErr': derived_params['apFluxErr'],
               'apMag': derived_params['apMag'], 'apMagErr': derived_params['apMagErr'],
               'trace': derived_params['trace'],
               'e1': derived_params['e1'], 'e2': derived_params['e2'],
               'e_final': derived_params['e_final'], 'phi_final': derived_params['phi_final'],
               'psf_fwhm': PSF_FWHM, 'objectId': objectId}
        return row

    def make_source_table_rowbyrow(self, save_file, method="analytical"):
        import time
        """
        Returns a source table generated from all the lens systems in the catalog
        under all the observation conditions in the observation history,
        and saves it as a .csv file.

        Keyword arguments:
        save_file -- path into which output source table will be saved
        method -- how to calculate moments for each row
                  (See method estimate_parameters for details about each option)
        """
        start = time.time()
        print("Began making the source catalog.")

        df = pd.DataFrame(columns=self.source_columns)
        #ellipticity_upper_limit = desc.slrealizer.get_ellipticity_cut()
        print("Number of systems: %d, number of observations: %d" %(self.num_systems, self.num_obs))

        hsm_failed = 0

        for j in xrange(self.num_obs):
            for i in xrange(self.num_systems):
                row = self.create_source_row(lens_info=self.get_lens_info(rownum=i),
                                             obs_info=self.observation.loc[j],
                                             method=method)
                if row == None:
                    hsm_failed += 1
                else:
                    df = df.append(row, ignore_index=True)

        df = df.infer_objects()
        df = df[self.source_columns]
        df.set_index('objectId', inplace=True)
        df.to_csv(save_file, index=True)

        end = time.time()
        if method == 'hsm':
            print("Done making the source table which has %d row(s) in %0.2f hours, after getting %d errors from HSM failure." %(len(df), (end - start)/3600.0, hsm_failed))
        else:
            print("Done making the source table with %s method in %0.2f minutes." %(method, (end - start)/60.0))
#        desc.slrealizer.dropbox_upload(dir, 'source_catalog_new.csv')

        self.sourceTable = df
        if self.DEBUG:
            return df

    def make_object_table(self, object_table_path, source_table_path=None, include_std=False):

        """
        Generates the object table from the given source table at source_table_path
        by averaging the properties for each filter, and saves it as object_table_path.
        """
        import time
        import gc

        if object_table_path is None:
            raise ValueError("Must provide save path of the output object table.")

        if source_table_path is not None:
            print("Reading in the source table at %s ..." %source_table_path)
            obj = pd.read_csv(source_table_path)
            obj.set_index('objectId', inplace=True)
        elif self.sourceTable is not None:
            print("Reading in Pandas Dataframe of most recent source table generated... ")
            obj = self.sourceTable.copy()
        else:
            raise ValueError("Must provide a source table path or generate a source table at least once using this Realizer object.")

        start = time.time()

        obj.drop(['ccdVisitId', 'psf_fwhm'], axis=1, inplace=True)
        # Define (filter-nonspecific) properties to go in object table columns
        keepCols = list(obj.columns.values)
        keepCols.remove('filter')
        # Pivot the filter values to the column axis
        obj = obj.pivot_table(values=keepCols, index=[obj.index, 'MJD'], columns='filter', aggfunc='first')
        # Collapse multi-indexed column using filter_property formatting
        obj.columns = obj.columns.map('{0[1]}_{0[0]}'.format)
        gc.collect()

        #if self.DEBUG: print(obj.columns.values)

        # Take mean, optional std of properties across observed times for each object
        obj.reset_index(inplace=True)
        obj.drop('MJD', axis=1, inplace=True)
        obj = obj.groupby('objectId', sort=False)
        means = obj.mean()
        if include_std:
            stds = obj.std()
            obj = means.join(stds, lsuffix='', rsuffix='-std')
        else:
            obj = means
        gc.collect()

        # Drop examples with missing values
        obj.dropna(how='any', inplace=True)
        # Get x, y values relative to the r-band
        for b in 'ugriz':
            obj[b + '_' + 'x'] = obj[b + '_' + 'x'] - obj['r_x']
            obj[b + '_' + 'y'] = obj[b + '_' + 'y'] - obj['r_y']
        end = time.time()

        # Save as csv file
        obj.to_csv(object_table_path, index=False)
        print("Done making the object table in %0.2f seconds." %(end-start))
        #if self.DEBUG:
            #print("Object table columns: ", obj.columns)

        #desc.slrealizer.dropbox_upload(save_dir, 'object_catalog_new.csv') #this uploads to the desc account

    def include_quasar_variability(self, save_output=False, input_source_path=None, output_source_path=None):
        """
        Takes a source table and adds the intrinsic variability of the quasar images
        using the generative model introduced in MacLeod et al (2010)

        Keyword arguments:
        save_output -- whether to save the output to disk [default: False]
        input_source_path -- path of input source table to be altered [default: None]
        output_source_path -- path of output source table containing time variability [default: None]
        """

        import gc
        import time

        start = time.time()
        if input_source_path is None:
            try:
                print("Reading in the most recent source table...")
                src = self.source_table
            except ValueError:
                print("Realizer has not generated a source table yet.")
        else:
            try:
                print("Reading in the source table at %s" %input_source_path)
                src = pd.read_csv(input_source_path)
            except ValueError:
                print("Please input a valid path to the source table.")

        if self.DEBUG:
            NUM_TIMES = src['MJD'].nunique()
            NUM_OBJECTS = src['objectId'].nunique()
            print(NUM_TIMES, NUM_OBJECTS)

        for q in range(4):
            magnitude_type = 'q_mag_' + str(q)
            # Filter-specific dictionary of apMag values to replace src apMag
            apMag_filter = {}
            for b in 'ugriz':
                #########################
                # Adding useful columns #
                #########################
                # Filter-specific version of src with only the columns we need
                src_timevar = src.query('filter == @b')[['filter', 'objectId', 'MJD', magnitude_type, ]].copy()
                # Total number of time steps in this filter
                NUM_TIMES_FILTER = src_timevar['MJD'].nunique()
                # Set MJD relative to zero
                src_timevar['MJD_diff'] = src_timevar['MJD'] - np.min(src_timevar['MJD'].values)
                # Map ccdVisitId to integers starting from 0
                sorted_obs_id = np.sort(src_timevar['MJD'].unique())
                time_map = dict(zip(sorted_obs_id, range(NUM_TIMES_FILTER)))
                src_timevar['time_index'] = src_timevar['MJD'].map(time_map).astype(int)
                # Add a column of time elapsed since last observation, d_time
                src_timevar.sort_values(['objectId', 'filter', 'MJD', 'MJD_diff'], axis=0, inplace=True)
                src_timevar['d_time'] = src_timevar['MJD_diff'] - src_timevar['MJD_diff'].shift(+1)
                src_timevar['d_time'].fillna(0.0, inplace=True)
                src_timevar['d_time'] = np.clip(src_timevar['d_time'], a_min=0.0, a_max=None)
                gc.collect()

                ##############################
                # Computing time variability #
                ##############################
                # Parameters of the generative model (hand-picked)
                MU = 0.0
                TAU = 20.0 #np.power(10.0, 2.4) # days
                S_INF = 0.14 # mag
                # Add a column to store the variability and initialize to zero
                src_timevar['intrinsic_mag'] = 0.0
                # Pivot to get time sequence to be horizontal
                src_timepivot = src_timevar.pivot_table(index=['filter', 'objectId'],
                                                        columns=['time_index',],
                                                        values=['d_time', magnitude_type, 'intrinsic_mag', ])
                # Update 'intrinsic_mag' column
                for t in range(1, NUM_TIMES_FILTER):
                    src_timepivot.loc[b, :]['intrinsic_mag'][t] = np.random.normal(loc=src_timepivot['intrinsic_mag'][t - 1].values*np.exp(-src_timepivot['d_time'][t].values/TAU) + MU*(1.0 - np.exp(-src_timepivot['d_time'][t].values/TAU)),
                                                                                scale=0.5*S_INF**2.0*(1.0 - np.exp(-2.0*src_timepivot['d_time'][t].values/TAU)))

                ########################
                # Final column sorting #
                ########################
                # Add computed variability to magnitude_type column
                src_timepivot[magnitude_type] = src_timepivot[magnitude_type] + src_timepivot['intrinsic_mag']
                # Pivot the time sequence back, to be vertical
                src_timepivot = src_timepivot.stack(dropna=False)
                # Drop columns we'll no longer use
                src_timepivot.drop(['intrinsic_mag', 'd_time'], axis=1, inplace=True)
                # As precaution, sort before storing updated apMag
                src_timepivot.reset_index(inplace=True)
                src_timepivot.sort_values(by=['filter', 'objectId', 'time_index'], inplace=True)
                # Store filter-specific apMag values
                apMag_filter[b] = src_timepivot[magnitude_type].values
                gc.collect()

            # Update apMag in source table, filter by filter
            for b in 'ugriz':
                src.loc[src['filter']==b, magnitude_type] = apMag_filter[b]

        gc.collect()
        if self.DEBUG:
            print("Result of adding time variability: ")
            print("Number of observations: ", src['MJD'].nunique())
            print("Number of objects: ", src['objectId'].nunique())

        src.set_index('objectId', inplace=True)
        end = time.time()

        print("Done adding time variability with %d row(s) in %0.2f seconds using vectorization." %(len(src), end-start))
        if save_output:
            print("Saving the new source table with time variability at %s" %output_source_path)
            src.to_csv(output_source_path)

        self.source_table = src

    def _include_moments(self, inplace=True, input_dict=None):
        """
        Adds columns of first and second moments (analytically computed)
        to self.source_table (if inplace=True) or some other dictionary

        Keyword arguments:
        inplace -- whether to alter the self.source_table.
                   If input_dict is None, will be assumed to be True. [default: True]
        input_dict -- object of type dict that will be used to compute the moments.
                      If inplace is True, will be ignored. [default: None]

        Returns (only if inplace is False):
        a new dictionary that is a version of the original input_dict with
        moment-related columns added
        """

        return_dict = False
        if inplace or input_dict is None:
            src = self.source_table
        elif input_dict is not None:
            is_dictionary = (type(input_dict) is dict)
            valid_columns = ['lens_flux', 'apFlux', 'e', 'beta', 'psf_fwhm']
            valid_columns += [value + '_%d' %idx for value in ['q_flux', 'XIMG', 'YIMG'] for idx in range(4)]
            has_valid_columns = all(col in input_dict for col in valid_columns)
            if is_dictionary and has_valid_columns:
               src = input_dict
               return_dict = True
            else:
                raise ValueError("Keyword input_dict must be a dictionary and contain all the required keys.")

        # Calculate flux ratios (for weighted moments)
        src['lensFluxRatio'] = src['lens_flux']/src['apFlux']
        for q in range(4):
            src['qFluxRatio_' + str(q)] = src['q_flux_' + str(q)]/src['apFlux']
        if not return_dict:
            src.drop(['lens_flux'] + ['q_flux_' + str(q) for q in range(4)], axis=1, inplace=True)

        #################
        # FIRST MOMENTS #
        #################
        src['x'], src['y'] = 0.0, 0.0
        for q in range(4):
            src['x'] += src['qFluxRatio_' + str(q)]*src['XIMG_' + str(q)]
            src['y'] += src['qFluxRatio_' + str(q)]*src['YIMG_' + str(q)]
        if self.add_moment_noise:
            src['x'] += utils.add_noise(mean=get_first_moment_err(),
                                  stdev=get_first_moment_err_std(),
                                  shape=src['x'].shape,
                                  measurement=src['x'])
            src['y'] += utils.add_noise(mean=get_first_moment_err(),
                                  stdev=get_first_moment_err_std(),
                                  shape=src['y'].shape,
                                  measurement=src['y'])
        ##################
        # SECOND MOMENTS #
        ##################
        # Read in lens shape parameters
        src['minor_to_major'] = np.power((1.0 - src['e'])/(1.0 + src['e']), 0.5) # q parameter in galsim.shear
        src['beta'] = np.radians(src['beta']) # beta parameter in galsim.shear
         # Arbitrarily set REFF_T to 1.0
        #src['sigmasq_lens'] = np.power(utils.hlr_to_sigma(src['REFF_T']), 2.0)
        src['sigmasq_lens'] = np.power(utils.hlr_to_sigma(1.0), 2.0)

        # Initialize with lens contributions
        src['lam1'] = src['sigmasq_lens']/src['minor_to_major']
        src['lam2'] = src['sigmasq_lens']*src['minor_to_major']
        src['lens_Ixx'] = src['lam1']*np.power(np.cos(src['beta']), 2.0) + src['lam2']*np.power(np.sin(src['beta']), 2.0)
        src['lens_Iyy'] = src['lam1']*np.power(np.sin(src['beta']), 2.0) + src['lam2']*np.power(np.cos(src['beta']), 2.0)
        src['lens_Ixy'] = (src['lam1'] - src['lam2'])*np.cos(src['beta'])*np.sin(src['beta'])
        src['Ixx'] = src['lensFluxRatio']*(src['lens_Ixx'] + np.power(src['x'], 2.0))
        src['Iyy'] = src['lensFluxRatio']*(src['lens_Iyy'] + np.power(src['y'], 2.0))
        src['Ixy'] = src['lensFluxRatio']*(src['lens_Ixy'] - src['x']*src['y'])
        if not return_dict:
            src.drop(['lam1', 'lam2', 'lens_Ixx', 'lens_Iyy', 'lens_Ixy'], axis=1, inplace=True)
        # Add quasar contributions
        for q in range(4):
            src['Ixx'] += src['qFluxRatio_' + str(q)]*(np.power(src['XIMG_' + str(q)] - src['x'], 2.0))
            src['Iyy'] += src['qFluxRatio_' + str(q)]*(np.power(src['YIMG_' + str(q)] - src['y'], 2.0))
            src['Ixy'] += src['qFluxRatio_' + str(q)]*(src['XIMG_' + str(q)] - src['x'])\
                                                     *(src['YIMG_' + str(q)] - src['y'])
        # Add PSF
        src['sigmasq_psf'] = np.power(utils.fwhm_to_sigma(src['psf_fwhm']), 2.0)
        src['Ixx'] += src['sigmasq_psf']
        src['Iyy'] += src['sigmasq_psf']

        # Get trace and ellipticities
        src['trace'] = src['Ixx'] + src['Iyy']
        if self.add_moment_noise:
            src['trace'] += utils.add_noise(mean=get_second_moment_err(),
                                      stdev=get_second_moment_err_std(),
                                      shape=src['trace'].shape,
                                      measurement=src['trace'])
        src['e1'] = (src['Ixx'] - src['Iyy'])/src['trace']
        src['e2'] = 2.0*src['Ixy']/src['trace']
        src['e_final'], src['phi_final'] = utils.e1e2_to_ephi(src['e1'], src['e2'])

        if inplace:
            self.source_table = src

        if return_dict:
            return src

    def compare_truth_vs_emulated(self, lensID=None, rownum=None, save_dir=None):
        """
        Draws two images of the lens system with the given rownum
        under randomly chosen observation conditions,
        one from the catalog info (truth image) and
        another from HSM's estimation of the truth image

        Keyword arguments:
        - lensID: an integer specifying the ID of the lens system in the OM10 catalog

        Returns:
        A tuple of
        - the truth image drawn from the catalog info
        - the emulated image drawn from HSM's derived info of aggregate system

        """
        # Only works in debug mode.
        self.debug = True
        # Randomly select observation ID
        obs_rownum = random.randint(0, self.num_obs)

        # Render the truth image
        truth_img = self.draw_system(self.get_lens_info(lensID=lensID, rownum=rownum),\
                                    self.get_observation(rownum=obs_rownum),\
                                    save_dir)

        # Render the emulated image under observation conditions indexed by obs_rownum
        fig, axes = plt.subplots(2, figsize=(5, 10))
        axes[0].imshow(truth_img.array, interpolation='none', aspect='auto')
        axes[0].set_title("TRUE MODEL IMAGE")
        estimated_params = self.estimate_hsm(image=truth_img, observation=self.observation.loc[obs_rownum])
        emulated_img = self.draw_emulated_system(estimated_params)
        axes[1].imshow(img.array, interpolation='none', aspect='auto')
        axes[1].set_title("EMULATED IMAGE")
        if save_dir is not None:
            plt.savefig(save_dir + 'truth_vs_emulated.png')
        return truth_img, emulated_img


import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
import cv2
import os
import time
from tensorflow.python.platform import gfile

# helper function to process image to input to graph
def preprocess_image(image, image_size):
    image_height, image_width = image.shape[:2]
    if image_height > image_width:
        scale = image_size / image_height
        resized_height = image_size
        resized_width = int(image_width * scale)
    else:
        scale = image_size / image_width
        resized_height = int(image_height * scale)
        resized_width = image_size

    image = cv2.resize(image, (resized_width, resized_height))
    new_image = np.ones((image_size, image_size, 3), dtype=np.float32) * 128.
    new_image[0:resized_height, 0:resized_width] = image.astype(np.float32)
    new_image /= 255.

    mean = [0.485, 0.456, 0.406]
    std = [0.229, 0.224, 0.225]

    for i in range(3):
        new_image[..., i] -= mean[i]
        new_image[..., i] /= std[i]

    return new_image, scale

# function to inference on videos, uses pb frozen graph
def inference_vid(model_path, phi=0, weighted_bifpn=False, classes=None, score_threshold=0.5, draw=True, fps=30):
  image_sizes = (512, 640, 768, 896, 1024, 1280, 1408)
  cv2.namedWindow('image',cv2.WINDOW_NORMAL)
  cv2.resizeWindow('image', 1000, 800)
  image_size = image_sizes[phi]
  num_classes = len(classes)
  colors = [np.random.randint(0, 256, 3).tolist() for _ in range(num_classes)]  #yellow

  tf.reset_default_graph()
  sess = tf.Session()
  with tf.gfile.GFile(model_path, 'rb') as f:
      graph_def = tf.GraphDef()
      graph_def.ParseFromString(f.read())
      sess.graph.as_default()
      tf.import_graph_def(graph_def, name='')

  tensor_input = sess.graph.get_tensor_by_name('input_1:0')

  output_boxes = sess.graph.get_tensor_by_name('filtered_detections/map/TensorArrayStack/TensorArrayGatherV3:0')
  output_scores = sess.graph.get_tensor_by_name('filtered_detections/map/TensorArrayStack_1/TensorArrayGatherV3:0')
  output_labels = sess.graph.get_tensor_by_name('filtered_detections/map/TensorArrayStack_2/TensorArrayGatherV3:0')

  cap = cv2.VideoCapture(0)
  while cap.isOpened():
    start = time.time()
    flag, image = cap.read()
    if flag:
        src_image = image.copy()
    else:
        break
    image = image[:, :, ::-1]   #convert from bgr to rgb
    h, w = image.shape[:2]

    image, scale = preprocess_image(image, image_size=image_size)
    # run network
    boxes, scores, labels = sess.run([output_boxes, output_scores, output_labels], {tensor_input:[image]})
    boxes /= scale
    boxes[0, :, 0] = np.clip(boxes[0, :, 0], 0, w - 1)
    boxes[0, :, 1] = np.clip(boxes[0, :, 1], 0, h - 1)
    boxes[0, :, 2] = np.clip(boxes[0, :, 2], 0, w - 1)
    boxes[0, :, 3] = np.clip(boxes[0, :, 3], 0, h - 1)

    # select indices which have a score above the threshold
    indices = np.where(scores[0, :] > score_threshold)[0]

    # select those detections
    boxes = boxes[0, indices]
    scores = scores[0, indices]
    labels = labels[0, indices]

    # draw the box
    for box, score, label in zip(boxes, scores, labels):
        xmin = int(round(box[0]))
        ymin = int(round(box[1]))
        xmax = int(round(box[2]))
        ymax = int(round(box[3]))
        score = '{:.4f}'.format(score)
        class_id = int(label)
        color = colors[class_id]
        class_name = classes[class_id]
        label = '-'.join([class_name, score])
        ret, baseline = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
        cv2.rectangle(src_image, (xmin, ymin), (xmax, ymax), color, 3)
        cv2.rectangle(src_image, (xmin, ymax - ret[1] - baseline), (xmin + ret[0], ymax), color, -1)
        cv2.putText(src_image, label, (xmin, ymax - baseline), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)
    # show camera feed
    cv2.imshow("image", src_image)
    cv2.waitKey(1)

if __name__ =='__main__':
    inference_vid('/home/craig/pascal_05.pb', phi=0, classes = [
        'aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair',
        'cow', 'diningtable', 'dog', 'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train',
        'tvmonitor',
    ])


from PIL import ImageGrab,ImageOps
from pynput.mouse import Button, Controller
import os,time
from numpy import *
from coordinates import *
# --------- Points -----------
"""
Set x_pad y_pad whit your game screem zero point, run getCords() in left-up coner.
Run screenGrab() and look if image is correctly adjusted.
"""
x_pad = 363
y_pad = 206

# Mause Comands

def leftClick():
    Controller().press(Button.left)
    Controller().release(Button.left)
    #print ("Click.")          #optional, for debugging purposes.

def mousePos(cord):
    Controller().position = ((x_pad + cord[0], y_pad + cord[1]))

# ---- Capture Tools

def screenGrab(x_pad,y_pad):
    """
    Return playable area.
    """
    box = (x_pad+1,y_pad+1,x_pad+637,y_pad+477)
    im = ImageGrab.grab(box)
    #im.save('full_snap__' + str(int(time.time())) + '.png', format='PNG')
    return im

def clear():
    """
    Clear all captured images in main code folder.
    """
    files = os.listdir()
    for file in files:
        if 'snap' in file:
            os.remove(file)
    print('Trash files are clear')

def getCords():
    """
    Get actual mouse location based in relative game window.
    """
    x,y = Controller().position         #
    x = x - x_pad
    y = y - y_pad
    print(x,y)
    return (x,y)

def capture():
    """
    Capture multiple cords in same time.
    """
    for i in range(6):
        time.sleep(1.5)
        getCords()

def status(item):
    """
    Return if a item is avalieble or not.

    Itens list:
    - shrimp
    - unagi
    - nori
    - roe
    - salmon
    - rice
    - sake
    """
    sg = screenGrab(x_pad,y_pad)
    cords, avaib = Cord.foodAvailabilityLocation, Cord.foodAvailability
    pixel = sg.getpixel(cords[item])
    if pixel == avaib[item]:
        print(f"{item} is avalible to buy")
        return True
    else:
        print(f"{item} is NOT avalible to buy")
        return False

def getPixel(tuple):
    """
    Return color of pixel based in x,y.
    """
    sg = screenGrab(x_pad,y_pad)
    print(sg.getpixel(tuple))
    return sg.getpixel(tuple)

def getSeat(table):
    """
    Return cliet order ID
    """
    leftCorn = list(Cord.clientOrder[table])
    leftCorn[0] += x_pad
    leftCorn[1] += y_pad
    box = (leftCorn[0],leftCorn[1],leftCorn[0]+58,leftCorn[1]+15)
    im = ImageOps.grayscale(ImageGrab.grab(box))
    a = array(im.getcolors())
    a = a.sum()
    #im.show()    # Debug
    #print(f"Id of order in {table} is: {a}")
    return a

# ---- Start Game

def startGame():
    """
    Collection of items to start the game.
    """
    #location of start button
    mousePos((310, 202))
    leftClick()
    time.sleep(.1)

    #location of skip donwload app
    mousePos((309, 389))
    leftClick()
    time.sleep(.1)

    #location of skip button
    mousePos((580, 450))
    leftClick()
    time.sleep(.1)

    #location of continue button
    mousePos((310, 372))
    leftClick()
    time.sleep(.1)

# ---- Cooking

def makeFood(food):
    """
    Recipes onigiri, californiaRoll, gunkaMaki
    """
    def rollSundari():
        mousePos(Cord.foodCord['sundari'])
        leftClick()
        time.sleep(.8)

    if food == 'onigiri':
        """
        Recipe: 2 rice 1 nori.
        """
        mousePos(Cord.foodCord['rice'])
        leftClick()
        time.sleep(.1)
        mousePos(Cord.foodCord['rice'])
        leftClick()
        time.sleep(.1)
        mousePos(Cord.foodCord['nori'])
        leftClick()
        Cord.stok['rice'] -= 2
        Cord.stok['nori'] -= 1

    elif food == 'californiaRoll':
        """
        Recipe: 1 rice 1 nori 1 roe.
        """
        mousePos(Cord.foodCord['rice'])
        leftClick()
        time.sleep(.1)
        mousePos(Cord.foodCord['nori'])
        leftClick()
        time.sleep(.1)
        mousePos(Cord.foodCord['roe'])
        leftClick()
        Cord.stok['rice'] -= 1
        Cord.stok['nori'] -= 1
        Cord.stok['roe'] -= 1

    elif food == 'gunkanMaki':
        """
        Recipe: 1 rice 1 nori 2 roe.
        """
        mousePos(Cord.foodCord['rice'])
        leftClick()
        time.sleep(.1)
        mousePos(Cord.foodCord['nori'])
        leftClick()
        time.sleep(.1)
        mousePos(Cord.foodCord['roe'])
        leftClick()
        time.sleep(.1)
        mousePos(Cord.foodCord['roe'])
        leftClick()
        Cord.stok['rice'] -= 1
        Cord.stok['nori'] -= 1
        Cord.stok['roe'] -= 2

    time.sleep(.1)
    rollSundari()
    time.sleep(1)
    print(f'make {food}') # debug

def clearTable(table):
    """
    Clear table based in table position x,y.
    """
    plate = Cord.plates[table]
    mousePos(plate)
    time.sleep(.1)
    leftClick()
    print(f'{table} are clear')

def buy(item):
    """
    Buy itens in shop fone.

    Itens list:
    - shrimp
    - unagi
    - nori
    - roe
    - salmon
    - rice
    - sake
    """
    cords = Cord.phoneMenu
    mousePos(cords["menuPhone"])
    leftClick()
    time.sleep(.1)
    if item in ("shrimp","unagi","nori","roe","salmon"):
        mousePos(cords["menuToppings"]["toppingsButton"])
        leftClick()
        time.sleep(.3)
    elif item == "rice":
        mousePos(cords["menuRice"]["riceButton"])
        leftClick()
        time.sleep(.3)
    elif item == "sake":
        mousePos(cords["menuSake"]["sakeButton"])
        leftClick()
        time.sleep(.3)
    if status(item):
        if item in ("shrimp","unagi","nori","roe","salmon"):
            mousePos(cords["menuToppings"][item])
            leftClick()
            time.sleep(.1)
            if item in ("shrimp","unagi"):
                Cord.stok[item] += 5
                Cord.stok["money"] -= 350
            if item == "nori":
                Cord.stok["nori"] += 10
                Cord.stok["money"] -= 100
            if item == "roe":
                Cord.stok["roe"] += 10
                Cord.stok["money"] -= 200
            if item == "salmon":
                Cord.stok["salmon"] += 5
                Cord.stok["money"] -= 300
        elif item == "rice":
            mousePos(cords["menuRice"]["rice"])
            leftClick()
            time.sleep(.1)
            Cord.stok["rice"] += 10
            Cord.stok["money"] -= 100
        elif item == "sake":
            mousePos(cords["menuSake"]["sake"])
            leftClick()
            time.sleep(.1)
            Cord.stok["salmon"] += 2
            Cord.stok["money"] -= 100
        mousePos(cords["menuDelivery"]["normal"])
        leftClick()
        time.sleep(1)
        return True
    else:
        print(f'{item} is not avaliable to buy')
        if item in ("shrimp","unagi","nori","roe","salmon"):
            mousePos(cords["menuToppings"]["exit"])
            leftClick()
            time.sleep(.1)
        else:
            mousePos(cords["menuRice"]["exit"])
            leftClick()
            time.sleep(.1)
        return False

def checkStok():
    """
    Check if stok is low and if's buy that item.
    """
    for food in Cord.stok:
        print(f"{food} has {Cord.stok[food]} left")
        if Cord.stok[food] <= 4:
            if food == "money":
                pass
            elif food == "sake":
                pass
            else:
                buy(food)
                print(f'Buy {food}')

def main():
    startGame()
    time.sleep(1)
    while True:
        for table in Cord.clientOrder:
            order = getSeat(table)
            if order in (8587,7442,3930):
                if not Cord.orders[table]:
                    checkStok()
                    makeFood(Cord.foodId[order])
                    Cord.orders[table] = True
            else:
                print (f"{table} unoccupied")
                clearTable(table)
                Cord.orders[table] = False
        time.sleep(1)

main()


import os

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"   # see issue #152
os.environ["CUDA_VISIBLE_DEVICES"] = ""

import pandas as pd
import numpy as np
import time
import functools
import pickle
from functools import partial

import keras
from keras.layers import Dense, Dropout, Activation
from keras.models import Sequential, load_model
from keras.optimizers import Adam
from keras.losses import categorical_crossentropy
from keras.metrics import top_k_categorical_accuracy
from keras.utils import Sequence
from keras.callbacks import EarlyStopping, CSVLogger, ModelCheckpoint, Callback, TensorBoard, ReduceLROnPlateau

from sklearn.preprocessing import LabelEncoder
from sklearn.utils import shuffle

from rdkit import Chem
from rdkit.Chem import AllChem, rdChemReactions, rdmolfiles, rdmolops, rdMolDescriptors
from rdkit.DataStructs import cDataStructs
from rdkit.Chem.Draw import IPythonConsole #Needed to show molecules
from rdkit.Chem.Draw.MolDrawing import MolDrawing, DrawingOptions #Only needed if modifying defaults

import sys
from rdchiral import main as rdc

class Model:
    """
    Class to facilate the predicition of synthetic steps for ring systems.
    Can only be used for the predicition of single synthetic steps for a given ring system.
    """
    def __init__(self, dataset="uspto_ringbreaker", mask=True):
        directory = os.getcwd()
        self.models = {"uspto_ringbreaker": os.path.join(directory, "models/checkpoints/weights.hdf5"),
                       "uspto_standard": os.path.join(directory,"models/uspto_standard/weights.hdf5"),
        }
        self.templates = {"uspto_ringbreaker": os.path.join(directory,"data/uspto_ringformations.csv"),
                          "uspto_standard": os.path.join(directory,"data/standard_uspto_template_library.csv"),
        }
        if 'filtered' in dataset:
            dataset = '_'.join(dataset.split('_')[:2])
            mask = True

        self.model = self.models[dataset]
        self.templates = self.templates[dataset]
        self.mask = mask

         #Initalise template library
        if 'ringbreaker' in dataset:
            lib = pd.read_csv(self.templates)
            lib = lib.drop(["Unnamed: 0", "index", "selectivity", "outcomes", "ring_change"], axis=1)
            template_labels = LabelEncoder()
            lib['template_code'] = template_labels.fit_transform(lib['template_hash'])
            lib = lib.drop(["reaction_hash", "reactants"], axis=1)
            self.lib = lib.set_index('template_code').T.to_dict('list')
            self.mask = False
        else:
            lib = pd.read_csv(self.templates, index_col=0, header=None, names=["index", "ID", "reaction_hash", "reactants", "products", "classification", "retro_template", "template_hash", "selectivity", "outcomes", "template_code"])
            lib = lib.drop(["reaction_hash", "reactants", "selectivity", "outcomes"], axis=1)
            self.lib = lib.set_index('template_code').T.to_dict('list')

        if mask == True and 'uspto' in dataset:
            with open("../data/uspto_filter_array.pkl", "rb") as f:
                self.filter_array = pickle.load(f)
        else:
            self.filter_array = None

        #Initialise policy
        # Define top k accuracies
        top10_acc = functools.partial(keras.metrics.top_k_categorical_accuracy, k=10)
        top10_acc.__name__ = 'top10_acc'

        top50_acc = functools.partial(keras.metrics.top_k_categorical_accuracy, k=50)
        top50_acc.__name__ = 'top50_acc'

        self.policy = load_model(self.model, custom_objects={'top10_acc': top10_acc, 'top50_acc': top50_acc})

    def smiles_to_ecfp(self, product, size=2048):
        """Converts a single SMILES into an ECFP4

        Parameters:
            product (str): The SMILES string corresponing to the product/molecule of choice.
            size (int): Size (dimensions) of the ECFP4 vector to be calculated.

        Returns:
            ecfp4 (arr): An n dimensional ECFP4 vector of the molecule.
        """
        mol = Chem.MolFromSmiles(product)
        ecfp = AllChem.GetMorganFingerprintAsBitVect(mol, 2, nBits = 2048)
        arr = np.zeros((0,), dtype=np.int8)
        cDataStructs.ConvertToNumpyArray(ecfp, arr)
        return arr.reshape([1,2048])

    def get_templates(self, prediction, lib, topN=50):
        """Given a prediction and the template library, obtains the top N predicted templates

        Parameters:
            prediction (np.array): The prediction from the softmax layer as a numpy array
            lib (dict): The template library as a dictionary where the template hash is the key

        Returns:
            predicted_templates (dict): The predicted templates and associated meta-data
        """
        sort_pred = np.argsort(prediction)[::-1]

        predicted_templates = {}
        for i in range(1, topN+1):
            pred_temp = lib[sort_pred[-1][-i]]
            predicted_templates[i] = {'template': pred_temp[-2], 'classification': pred_temp[-3], 'ID': pred_temp[0]}

        return predicted_templates

    def num_rings(self, reaction):
        """Given a reaction SMILES, obtains the change in the number of rings

        Parameters:
            reaction (SMILES): The reaction without agents in the form of SMILES

        Returns:
            num_rings (int): The change in the number of rings in the reaction
        """
        mol_r = Chem.MolFromSmiles(reaction.split('>>')[0])
        mol_p = Chem.MolFromSmiles(reaction.split('>>')[-1])

        return rdMolDescriptors.CalcNumRings(mol_p) - rdMolDescriptors.CalcNumRings(mol_r)

    def get_prediction(self, target):
        """Given a prediction and the template library, obtains the top N predicted templates

        Parameters:
            reaction (SMILES): The reaction without agents in the form of SMILES

        Returns:
            num_rings (int): The change in the number of rings in the reaction
        """
        prediction = self.policy.predict(self.smiles_to_ecfp(target))
        if self.mask:
            prediction = np.multiply(prediction, self.filter_array)
            return prediction
        else:
            return prediction

    def predict_ring_outcomes(self, target, cutoff=50):
        """Given a SMILES predicts the top N ring forming reactions and returns the results as a dictionary

        Parameters:
            target (SMILES): The target SMILES
            cutoff (int): The Top N predictions

        Returns:
            results (dict): The results in the form of a dictionary, contains the precursors, probability, cumulative probability,
                            ID, and the number of changed rings
        """
        prediction = self.get_prediction(target)

        predicted_templates = self.get_templates(prediction, self.lib, cutoff+1)
        sort_pred = np.sort(prediction)[::-1]

        results = {"prediction": [],
                   "probability": [],
                   "cumulative_probability": [],
                   "id": [],
                   "precursor": [],
                   "ring_change": []
                  }

        num_outcomes = 0
        for i in range(1, cutoff+1):
            template = predicted_templates[i]['template']
            outcomes = rdc.rdchiralRunText(template, target)
            if len(outcomes) == 0:
                continue
            if self.num_rings('>>'.join([outcomes[0], target])) >= 1:
                results["prediction"].append(i)
                results["probability"].append(sort_pred[-1][-i])
                results["cumulative_probability"].append(sum(sort_pred[-1][-i:]))
                results["id"].append(predicted_templates[i]['ID'])
                results["precursor"].append(outcomes[0])
                results["ring_change"].append(self.num_rings('>>'.join([outcomes[0], target])))
                num_outcomes += 1
            else:
                continue
        results["outcomes"] = num_outcomes

        return results

    def show_topN_predictions(self, target, topN=50):
        """Given a SMILES predicts the top N ring forming reactions and displays the predictions for visualisation in Jupyter Notebook

        Parameters:
            target (SMILES): The target SMILES
        """
        prediction = self.get_prediction(target)

        predicted_templates = self.get_templates(prediction, self.lib, topN+1)
        sort_pred = np.sort(prediction)[::-1]

        for i in range(1, topN+1):
            template = predicted_templates[i]['template']
            outcomes = rdc.rdchiralRunText(template, target)
            if len(outcomes) != 0:
                print('###Prediction {}###'.format(str(i)))
                display(rdChemReactions.ReactionFromSmarts(template))
                print(sort_pred[-1][-i])
                print('Classification')
                print(predicted_templates[i]['classification'])
                print('Patent ID')
                print(predicted_templates[i]['ID'])
                print('---List of precursors---')
                print(outcomes)
                display(Chem.MolFromSmiles(outcomes[0]))
                print('\n')


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Implements several popular block and stream ciphers.
"""
from typing import Iterable, Any, ByteString
from Crypto.Util.Padding import pad, unpad
from Crypto.Util import Counter

from ... import arg, Unit, Executable, RefineryCriticalException, RefineryPartialResult
from ....lib.argformats import OptionFactory, extract_options

class CipherExecutable(Executable):
    """
    A metaclass for the abstract class `refinery.units.crypto.cipher.CipherUnit` which
    normalizes the class variable `key_sizes` containing an iterable of all possible
    key sizes that are acceptable for the represented cipher.
    """

    def __new__(mcs, name, bases, nmspc, abstract=False, blocksize=1, key_sizes=None):
        nmspc.setdefault('blocksize', blocksize)
        nmspc.setdefault('key_sizes', key_sizes)
        return super(CipherExecutable, mcs).__new__(mcs, name, bases, nmspc, abstract=abstract)

    def __init__(cls, name, bases, nmspc, abstract=False, **_):
        cls.key_sizes = (cls.key_sizes,) if isinstance(cls.key_sizes, int) else tuple(cls.key_sizes or ())
        super(CipherExecutable, cls).__init__(name, bases, nmspc, abstract=abstract)

class CipherUnit(Unit, metaclass=CipherExecutable, abstract=True):

    def __init__(self, key: arg(help='The encryption key.'), **keywords):
        super().__init__(key=key, **keywords)

    def decrypt(self, data: ByteString) -> ByteString:
        raise NotImplementedError

    def encrypt(self, data: ByteString) -> ByteString:
        raise NotImplementedError

    def process(self, data: ByteString) -> ByteString:
        if self.key_sizes and len(self.args.key) not in self.key_sizes:
            raise ValueError(F'the given key has an invalid length of {len(self.args.key)} bytes.')
        return self.decrypt(data)

    def reverse(self, data: ByteString) -> ByteString:
        return self.encrypt(data)

class StreamCipherUnit(CipherUnit, abstract=True):

    def __init__(self, key, **keywords):
        super().__init__(key=key, **keywords)

    def keystream(self) -> Iterable[int]:
        raise NotImplementedError

    def encrypt(self, data: bytearray) -> bytearray:
        import numpy as np
        key = np.fromiter(
            self.keystream(), dtype=np.uint8, count=len(data))
        out = np.frombuffer(
            memoryview(data), dtype=np.uint8, count=len(data))
        out ^= key
        return out

    decrypt = encrypt

class BlockCipherUnitBase(CipherUnit, abstract=True):
    def __init__(
        self, key, iv: arg('-I', '--iv', help=(
            'Specifies the initialization vector. If none is specified, then a block '
            'of zero bytes is used.')) = B'',
        padding: arg.choice('-P', choices=['PKCS7', 'ISO7816', 'X923', 'RAW'],
            nargs=1, metavar='ALG', help=(
            'Choose a padding algorithm ({choices}). The RAW algorithm does nothing. '
            'By default, all other algorithms are attempted. In most cases, the data '
            'was not correctly decrypted if none of these work.')
        ) = None,
        **keywords
    ):
        if not padding:
            padding = ['PKCS7', 'ISO7816', 'X923']
        elif not isinstance(padding, list):
            padding = [padding]
        iv = iv or bytes(self.blocksize)
        super().__init__(key=key, iv=iv, padding=padding, **keywords)

    @property
    def iv(self) -> ByteString:
        return self.args.iv

    def reverse(self, data: ByteString) -> ByteString:
        padding = self.args.padding[0]
        self.log_info('padding method:', padding)
        if padding != 'RAW':
            data = pad(data, self.blocksize, padding.lower())
        return super().reverse(data)

    def process(self, data: ByteString) -> ByteString:
        result = super().process(data)
        for p in self.args.padding:
            if p == 'RAW':
                return result
            try:
                return unpad(result, self.blocksize, p.lower())
            except Exception:
                pass
        raise RefineryPartialResult(
            'None of these paddings worked: {}'.format(', '.join(self.args.padding)),
            partial=result)

class StandardCipherExecutable(CipherExecutable):

    def __new__(mcs, name, bases, nmspc, cipher=NotImplemented):
        if cipher is NotImplemented:
            keywords = dict(abstract=True)
        else:
            keywords = dict(
                abstract=False,
                blocksize=cipher.block_size,
                key_sizes=cipher.key_size,
            )
        return super(StandardCipherExecutable, mcs).__new__(mcs, name, bases, nmspc, **keywords)

    def __init__(cls, name, bases, nmspc, cipher=NotImplemented):
        super(StandardCipherExecutable, cls).__init__(
            name, bases, nmspc, abstract=cipher is NotImplemented)
        cls._stdcipher_ = cipher
        if cipher is not NotImplemented and cipher.block_size > 1 and 'mode' in cls._argspec_:
            modes = extract_options(cipher)
            if not modes:
                raise RefineryCriticalException(
                    F'The cipher {cipher.name} is a block cipher module, '
                    F'but no cipher block mode constants were found.'
                )
            cls._bcmspec_ = OptionFactory(modes, ignorecase=True)
            cls._argspec_['mode'].merge_all(arg(
                '-M', '--mode', type=str.upper, metavar='MODE', nargs=arg.delete, choices=list(modes),
                help=(
                    'Choose cipher mode to be used. Possible values are: {}. By default, the CBC mode'
                    '  is used when an IV is is provided, and ECB otherwise.'.format(', '.join(modes))
                )
            ))

class StandardCipherUnit(CipherUnit, metaclass=StandardCipherExecutable):

    def __init__(self, key, **keywords):
        super().__init__(key, **keywords)

    def _get_cipher_instance(self, **optionals) -> Any:
        self.log_info(lambda: F'encryption key: {self.args.key.hex()}')
        return self._stdcipher_.new(key=self.args.key, **optionals)

    def encrypt(self, data: bytes) -> bytes:
        return self._get_cipher_instance().encrypt(data)

    def decrypt(self, data: bytes) -> bytes:
        return self._get_cipher_instance().decrypt(data)

class StandardBlockCipherUnit(BlockCipherUnitBase, StandardCipherUnit):

    def __init__(self, key, iv=B'', padding=None, mode=None):
        mode = self._bcmspec_(mode or iv and 'CBC' or 'ECB')
        if iv:
            if mode.name == 'ECB':
                raise ValueError('No initialization vector can be specified for ECB mode.')
            bounds = {
                'CCM': (7, 14),
                'OCB': (1, 16),
                'CTR': (1, 17),
            }.get(mode.name, None)
            if bounds and len(iv) not in range(*bounds):
                raise ValueError(F'Invalid nonce length, must be in {bounds} for {mode.name}.')
        super().__init__(key=key, iv=iv, padding=padding, mode=mode)

    def _get_cipher_instance(self, **optionals) -> Any:
        mode = self.args.mode.name
        if mode != 'ECB':
            if mode == 'CTR' and len(self.iv) == 16:
                counter = Counter.new(self.blocksize * 8,
                    initial_value=int.from_bytes(self.iv, 'big'))
                optionals['counter'] = counter
            elif mode in ('CCM', 'EAX', 'GCM', 'SIV', 'OCB', 'CTR'):
                optionals['nonce'] = self.iv
            elif mode in ('CBC', 'CFB', 'OFB', 'OPENPGP'):
                iv = self.iv
                if len(iv) != self.blocksize:
                    raise ValueError(F'The IV has length {len(self.args.iv)} but the block size is {self.blocksize}.')
                optionals['iv'] = self.iv
            self.log_info('initial vector:', self.iv.hex())
        if self.args.mode:
            optionals['mode'] = self.args.mode.value
        try:
            return super()._get_cipher_instance(**optionals)
        except TypeError:
            if 'iv' not in optionals:
                raise
            del optionals['iv']
            if self.iv:
                self.log_info('ignoring iv for mode', self.args.mode)
            return self._stdcipher_.new(key=self.args.key, **optionals)

class LatinCipherUnit(StreamCipherUnit, abstract=True):
    key_sizes = 16, 32

    def __init__(
        self, key,
        nonce: arg(help='The nonce. Default is the string {default}.') = B'REFINERY',
        magic: arg('-m', help='The magic constant; depends on the key size by default.') = B'',
        offset: arg.number('-x', help='Optionally specify the stream index, default is {default}.') = 0,
        rounds: arg.number('-r', help='The number of rounds. Has to be an even number.') = 20,
    ):
        super().__init__(key=key, nonce=nonce, magic=magic, offset=offset, rounds=rounds)

class LatinCipherStandardUnit(StandardCipherUnit):
    def __init__(self, key, nonce: arg(help='The nonce. Default is the string {default}.') = B'REFINERY'):
        super().__init__(key, nonce=nonce)

    def _get_cipher_instance(self, **optionals) -> Any:
        self.log_info('one-time nonce:', self.args.nonce.hex())
        return super()._get_cipher_instance(nonce=self.args.nonce)


import _thread
import argparse
import logging
import os
import random
import threading
from pathlib import Path

import numpy as np
import paddle
import soundfile
import yaml
from tqdm import tqdm
from yacs.config import CfgNode
from parakeet.models.fastspeech2 import FastSpeech2, FastSpeech2Inference
from parakeet.models.parallel_wavegan import PWGGenerator, PWGInference
from parakeet.modules.normalizer import ZScore

from frontend import Frontend

def generate(args, fastspeech2_config, pwg_config):
    # dataloader has been too verbose
    logging.getLogger("DataLoader").disabled = True

    # construct dataset for generate
    sentences = []
    with open(args.text, 'rt', encoding='utf-8') as f:
        for line in f:
            utt_id, sentence = line.strip().split()
            sentences.append((utt_id, sentence))

    with open(args.phones_dict, "r", encoding='utf-8') as f:
        phn_id = [line.strip().split() for line in f.readlines()]
    vocab_size = len(phn_id)
    print("vocab_size:", vocab_size)
    with open(args.speaker_dict, 'rt', encoding='utf-8') as f:
        spk_id = [line.strip().split() for line in f.readlines()]
    num_speakers = len(spk_id)
    print("num_speakers:", num_speakers)

    odim = fastspeech2_config.n_mels
    model = FastSpeech2(idim=vocab_size,
                        odim=odim,
                        num_speakers=num_speakers,
                        **fastspeech2_config["model"])

    model.set_state_dict(paddle.load(args.fastspeech2_checkpoint)["main_params"])
    model.eval()

    vocoder = PWGGenerator(**pwg_config["generator_params"])
    vocoder.set_state_dict(paddle.load(args.pwg_params))
    vocoder.remove_weight_norm()
    vocoder.eval()
    print("model done!")

    frontend = Frontend(args.phones_dict)
    print("frontend done!")

    stat = np.load(args.fastspeech2_stat)
    mu, std = stat
    mu = paddle.to_tensor(mu)
    std = paddle.to_tensor(std)
    fastspeech2_normalizer = ZScore(mu, std)

    stat = np.load(args.pwg_stat)
    mu, std = stat
    mu = paddle.to_tensor(mu)
    std = paddle.to_tensor(std)
    pwg_normalizer = ZScore(mu, std)

    fastspeech2_inference = FastSpeech2Inference(fastspeech2_normalizer, model)
    pwg_inference = PWGInference(pwg_normalizer, vocoder)

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    annotation_path = Path(os.path.dirname(args.annotation_path))
    annotation_path.mkdir(parents=True, exist_ok=True)
    start_num = 0
    if os.path.exists(args.annotation_path):
        with open(args.annotation_path, 'r', encoding='utf-8') as f_ann:
            start_num = len(f_ann.readlines())
    f_ann = open(args.annotation_path, 'a', encoding='utf-8')
    # 开始生成音频
    for i in tqdm(range(start_num, len(sentences))):
        utt_id, sentence = sentences[i]
        # 随机说话人
        spk_id = random.randint(0, num_speakers - 1)
        try:
            input_ids = frontend.get_input_ids(sentence, merge_sentences=True)
        except:
            continue
        phone_ids = input_ids["phone_ids"]
        flags = 0
        for part_phone_ids in phone_ids:
            with paddle.no_grad():
                mel = fastspeech2_inference(part_phone_ids, spk_id=paddle.to_tensor(spk_id))
                temp_wav = pwg_inference(mel)
                if flags == 0:
                    wav = temp_wav
                    flags = 1
                else:
                    wav = paddle.concat([wav, temp_wav])
        save_audio_path = str(output_dir / (utt_id + ".wav"))
        soundfile.write(save_audio_path, wav.numpy(), samplerate=fastspeech2_config.fs)
        f_ann.write('%s\t%s\n' % (save_audio_path[6:].replace('\\', '/'), sentence.replace('。', '').replace('，', '')
                                  .replace('！', '').replace('？', '')))
        f_ann.flush()

def main():
    parser = argparse.ArgumentParser(description="Synthesize with fastspeech2 & parallel wavegan.")
    parser.add_argument("--fastspeech2-config",
                        type=str,
                        default='models/fastspeech2_nosil_aishell3_ckpt_0.4/default.yaml',
                        help="fastspeech2 config file to overwrite default config.")
    parser.add_argument("--fastspeech2-checkpoint",
                        type=str,
                        default='models/fastspeech2_nosil_aishell3_ckpt_0.4/snapshot_iter_96400.pdz',
                        help="fastspeech2 checkpoint to load.")
    parser.add_argument("--fastspeech2-stat",
                        type=str,
                        default='models/fastspeech2_nosil_aishell3_ckpt_0.4/speech_stats.npy',
                        help="mean and standard deviation used to normalize spectrogram when training fastspeech2.")
    parser.add_argument("--pwg-config",
                        type=str,
                        default='models/parallel_wavegan_baker_ckpt_0.4/pwg_default.yaml',
                        help="parallel wavegan config file to overwrite default config.")
    parser.add_argument("--pwg-params",
                        type=str,
                        default='models/parallel_wavegan_baker_ckpt_0.4/pwg_generator.pdparams',
                        help="parallel wavegan generator parameters to load.")
    parser.add_argument("--pwg-stat",
                        type=str,
                        default='models/parallel_wavegan_baker_ckpt_0.4/pwg_stats.npy',
                        help="mean and standard deviation used to normalize spectrogram when training parallel wavegan.")
    parser.add_argument("--phones-dict",
                        type=str,
                        default="models/fastspeech2_nosil_aishell3_ckpt_0.4/phone_id_map.txt",
                        help="phone vocabulary file.")
    parser.add_argument("--speaker-dict",
                        type=str,
                        default="models/fastspeech2_nosil_aishell3_ckpt_0.4/speaker_id_map.txt",
                        help="speaker id map file.")
    parser.add_argument("--text",
                        type=str,
                        default='corpus.txt',
                        help="text to synthesize, a 'utt_id sentence' pair per line.")
    parser.add_argument("--output_dir", type=str, default='../../dataset/audio/generate', help="output audio dir.")
    parser.add_argument("--annotation_path", type=str, default='../../dataset/annotation/generate.txt',
                        help="audio annotation path.")
    parser.add_argument("--device", type=str, default="gpu", help="device type to use.")
    parser.add_argument("--verbose", type=int, default=1, help="verbose.")

    args = parser.parse_args()
    with open(args.fastspeech2_config) as f:
        fastspeech2_config = CfgNode(yaml.safe_load(f))
    with open(args.pwg_config) as f:
        pwg_config = CfgNode(yaml.safe_load(f))

    print("========Args========")
    print(yaml.safe_dump(vars(args)))
    print("========Config========")
    print(fastspeech2_config)
    print(pwg_config)

    generate(args, fastspeech2_config, pwg_config)

if __name__ == "__main__":
    main()


import numpy as np
from string import punctuation
from collections import defaultdict
import operator as op
from bisect import bisect
from pyteomics import parser, mass, fasta, auxiliary as aux, mgf, mzml
from . import scoring, utils
import logging
logger = logging.getLogger(__name__)

try:
    from pyteomics import cmass
except ImportError:
#   logger.warning('cmass could not be imported')
    cmass = mass

try:
    # import pyximport; pyximport.install()
    from .cutils import theor_spectrum
except:
    logger.info('Cython modules were not loaded...')
    from .utils import theor_spectrum
from .utils import reshape_theor_spectrum

spectra = {}
titles = {}
best_res = {}
nmasses = {}
t2s = {}
charges = {}
def prepare_peptide_processor(fname, settings):
    global spectra
    global nmasses
    global titles
    global t2s
    global charges
    global best_res
    best_res = {}
    maxcharges = {}
    fcharge = settings.getint('scoring', 'maximum fragment charge')
    ch_range = range(settings.getint('search', 'minimum charge'),
                1 + settings.getint('search', 'maximum charge'))
    for c in ch_range:
        maxcharges[c] = max(1, min(fcharge, c-1) if fcharge else c-1)

    params = {}
    params['maxpeaks'] = settings.getint('scoring', 'maximum peaks')
    params['minpeaks'] = settings.getint('scoring', 'minimum peaks')
    params['dynrange'] = settings.getfloat('scoring', 'dynamic range')
    params['acc'] = settings.getfloat('search', 'product accuracy')
    params['min_mz'] = settings.getfloat('search', 'product minimum m/z')
    params.update(utils._charge_params(settings))
    params['dacc'] = settings.getfloat('input', 'deisotoping mass tolerance')
    params['deisotope'] = settings.getboolean('input', 'deisotope')
    params['tags'] = utils.get_tags(settings.get('output', 'tags'))

    if not spectra:
        logger.info('Reading spectra ...')
        for spec in utils.iterate_spectra(fname):
            ps = utils.preprocess_spectrum(spec, params)
            if ps is not None:
                ttl = utils.get_title(ps)
                t2s[ttl] = ps
                for m, c in utils.neutral_masses(ps, params):
                    effc = maxcharges[c]
                    nmasses.setdefault(effc, []).append(m)
                    spectra.setdefault(effc, []).append(ps)
                    titles.setdefault(effc, []).append(ttl)
                    charges.setdefault(effc, []).append(c)
                    ps.setdefault('nm', {})[c] = m
        logger.info('%s spectra pass quality criteria.', sum(map(len, spectra.itervalues())))
        for c in list(spectra):
            i = np.argsort(nmasses[c])
            nmasses[c] = np.array(nmasses[c])[i]
            spectra[c] = np.array(spectra[c])[i]
            titles[c] = np.array(titles[c])[i]
            charges[c] = np.array(charges[c])[i]
    else:
        logger.info('Reusing %s spectra from previous run.', sum(map(len, spectra.itervalues())))

    utils.set_mod_dict(settings)

    mods = settings.get('modifications', 'variable')
    maxmods = settings.getint('modifications', 'maximum variable mods')
    leg = settings.get('misc', 'legend')
    punct = set(punctuation)
    nmods = [(p, mod[1], mod[2]) for p, mod in leg.iteritems() if p in punct]

    aa_mass = utils.get_aa_mass(settings)
    score = utils.import_(settings.get('scoring', 'score'))
    try:
        score_fast_name = settings.get('scoring', 'score') + '_fast'
        if score_fast_name == 'identipy.scoring.RNHS_fast':
            try:
                from cutils import RNHS_fast as score_fast
            except:
                score_fast = utils.import_(settings.get('scoring', 'score') + '_fast')
        else:
            score_fast = utils.import_(settings.get('scoring', 'score') + '_fast')
    except Exception as e:
        score_fast = False
        logging.debug('No fast score imported: %s', e)
    acc_l = settings.getfloat('search', 'precursor accuracy left')
    acc_r = settings.getfloat('search', 'precursor accuracy right')
    acc_frag = settings.getfloat('search', 'product accuracy')
    frag_unit = settings.get('search', 'product accuracy unit')
    if frag_unit == 'ppm':
        acc_frag_ppm = settings.getfloat('search', 'product accuracy ppm')
    else:
        acc_frag_ppm = False
    try:
        fast_first_stage = settings.getint('misc', 'fast first stage')
    except:
        fast_first_stage = 0
    unit = settings.get('search', 'precursor accuracy unit')
    rel = utils.relative(unit)

    if settings.has_option('scoring', 'condition'):
        cond = settings.get('scoring', 'condition')
    else:
        cond = None
    if isinstance(cond, str) and cond.strip():
        cond = utils.import_(cond)

    score = utils.import_(settings.get('scoring', 'score'))

    return {'rel': rel, 'aa_mass': aa_mass, 'acc_l': acc_l, 'acc_r': acc_r, 'acc_frag': acc_frag, 'acc_frag_ppm': acc_frag_ppm,
            'unit': unit, 'nmods': nmods, 'maxmods': maxmods, 'fast first stage': fast_first_stage,
            'sapime': utils.get_shifts_and_pime(settings),
            'cond': cond, 'score': score, 'score_fast': score_fast,
            'settings': settings}

def peptide_processor_iter_isoforms(peptide, **kwargs):
    nmods, maxmods = op.itemgetter('nmods', 'maxmods')(kwargs)
    if nmods and maxmods:
        out = []
        for form in utils.custom_isoforms(peptide, variable_mods=nmods, maxmods=maxmods, snp=kwargs['snp']):
            res = peptide_processor(form, **kwargs)
            if res:
                out.append(res)
        if out:
            return out
    else:
        res = peptide_processor(peptide, **kwargs)
        if res:
            return [res, ]

def peptide_processor(peptide, **kwargs):
    if kwargs['snp']:
        if 'snp' not in peptide:
            seqm = peptide
            aachange_pos = False
            snp_label = 'wild'
        else:
            tmp = peptide.split('snp')
            seqm = tmp[0] + tmp[1].split('at')[0].split('to')[-1] + tmp[2]
            aachange_pos = len(tmp[0]) + 1
            snp_label = tmp[1]
    else:
        seqm = peptide
        aachange_pos = False
        snp_label = False
    nterm_mass = kwargs.get('nterm_mass')
    cterm_mass = kwargs.get('cterm_mass')
    m = utils.custom_mass(seqm, aa_mass=kwargs['aa_mass'], nterm_mass = nterm_mass, cterm_mass = cterm_mass)
    # m = cmass.fast_mass(seqm, aa_mass=kwargs['aa_mass']) + (nterm_mass - 1.007825) + (cterm_mass - 17.002735)
    rel = kwargs['rel']
    acc_l = kwargs['acc_l']
    acc_r = kwargs['acc_r']
    settings = kwargs['settings']
    shifts_and_pime = kwargs['sapime']
    theor = {}
    theoretical_set = {}
    cand_idx = {}
    stored_value = False
    if rel:
        dm_l = acc_l * m / 1.0e6
        dm_r = acc_r * m / 1.0e6
    for c in spectra:

        if not rel:
            dm_l = acc_l * c
            dm_r = acc_r * c
        idx = set()
        for shift in shifts_and_pime:
            start = nmasses[c].searchsorted(m + shift - dm_l)
            end   = nmasses[c].searchsorted(m + shift + dm_r)
            if end - start:
                idx.update(range(start, end))
        if kwargs['cond']:
            idx2 = set()
            for i in idx:
                cond_val, stored_value = kwargs['cond'](spectra[c][i], seqm, settings, stored_value)
                if cond_val:
                    idx2.add(i)
            idx = idx2

        if idx:
            cand_idx[c] = idx
            theor[c], theoretical_set[c] = theor_spectrum(seqm, maxcharge=c, aa_mass=kwargs['aa_mass'], reshape=False,
                                                          acc_frag=kwargs['acc_frag'], nterm_mass = nterm_mass,
                                                          cterm_mass = cterm_mass, nm=m)
            reshaped = False

    results = []
    for fc, ind in cand_idx.iteritems():
        reshaped = False
        for i in ind:
            s = spectra[fc][i]
            # st = utils.get_title(s)
            st = titles[fc][i]
            if kwargs['score_fast']:
                hf = kwargs['score_fast'](s['fastset'], s['idict'], theoretical_set[fc], kwargs['min_matched'])
                if hf[0]:
                    if -hf[1] <= best_res.get(st, 0):
                        if kwargs['fast first stage']:
                            sc = hf[1]
                            score = {'match': [], 'sumI': 1, 'dist': [], 'total_matched': 999}
                        else:
                            if not reshaped:
                                theor[fc] = reshape_theor_spectrum(theor[fc])
                                reshaped = True
                            score = kwargs['score'](s, theor[fc], kwargs['acc_frag'], kwargs['acc_frag_ppm'], position=aachange_pos)#settings.getfloat('search', 'product accuracy ppm'))  # FIXME (?)
                            sc = score.pop('score')
                        if -sc <= best_res.get(st, 0) and score.pop('total_matched') >= kwargs['min_matched']:
                            results.append((sc, st, score, m, charges[fc][i], snp_label))
            else:
                # st = utils.get_title(s)
                if not reshaped:
                    theor[fc] = reshape_theor_spectrum(theor[fc])
                    reshaped = True
                score = kwargs['score'](s, theor[fc], kwargs['acc_frag'], kwargs['acc_frag_ppm'], position=aachange_pos)#settings.getfloat('search', 'product accuracy ppm'))  # FIXME (?)
                sc = score.pop('score')
                if -sc <= best_res.get(st, 0) and score.pop('total_matched') >= kwargs['min_matched']:
                    results.append((sc, st, score, m, charges[fc][i], snp_label))

    # results.sort(reverse=True, key=op.itemgetter(0))
    # results = np.array(results, dtype=[('score', np.float32), ('title', np.str_, 30), ('spectrum', np.object_), ('info', np.object_)])
    if results:
        return seqm, results
    # return seqm, []

def process_peptides(fname, settings):
    spec_results = defaultdict(dict)
    peps = utils.peptide_gen(settings)
    kwargs = prepare_peptide_processor(fname, settings)
    func = peptide_processor_iter_isoforms
    kwargs['min_matched'] = settings.getint('output', 'minimum matched')
    kwargs['snp'] = settings.getint('search', 'snp')
    kwargs['nterm_mass'] = settings.getfloat('modifications', 'protein nterm cleavage')
    kwargs['cterm_mass'] = settings.getfloat('modifications', 'protein cterm cleavage')
    kwargs['qsize'] = settings.getint('performance', 'out queue size')
    logger.info('Running the search ...')
    n = settings.getint('performance', 'processes')
    leg = {}
    if settings.has_option('misc', 'legend'):
        leg = settings.get('misc', 'legend')
    for y in utils.multimap(n, func, peps, **kwargs):
        for x in y:
            if x[1] is not None:
                peptide, result = x
                for score, spec_t, info, m, c, snp_label in result:
                    spec_results[spec_t]['spectrum'] = t2s[spec_t]
                    top_scores = spec_results[spec_t].setdefault('top_scores', 0)
                    if -score <= top_scores:
                        best_res[spec_t] = -score
                        info['pep_nm'] = m
                        info['charge'] = c
                        spec_results[spec_t]['top_scores'] = -score
                        spec_results[spec_t]['sequences'] = peptide
                        spec_results[spec_t]['info'] = info
                        spec_results[spec_t]['snp_label'] = snp_label

#               spec_results[spec_t].setdefault('scores', []).append(score) FIXME write histogram
#
#                     top_seqs   = spec_results[spec_t].setdefault('sequences', '')
#                     top_info   = spec_results[spec_t].setdefault('info', [])
#
#                     i = bisect(top_scores, -score)
#                     if nc is None or i < nc:
#                         top_scores.insert(i, -score)
#                         top_seqs.insert(i, peptide)
#                         top_info.insert(i, info)
#                         if nc is not None and len(top_scores) > nc:
#                             top_scores.pop()
#                             top_seqs.pop()
#                             top_info.pop()
    maxlen = settings.getint('search', 'peptide maximum length')
    dtype = np.dtype([('score', np.float64),
        ('seq', np.str_, maxlen), ('note', np.str_, 1),
        ('charge', np.int8), ('info', np.object_), ('sumI', np.float64), ('fragmentMT', np.float64), ('snp_label', np.str_, 15)])
    for spec_name, val in spec_results.iteritems():
        s = val['spectrum']
        c = []
        evalues = []
        score = val['top_scores']
        # for idx, score in enumerate(val['top_scores']):
        mseq = val['sequences']#[idx]
        seq = mseq
        info = val['info']#[idx]
        for x in set(mseq).intersection(punctuation):
            repl = leg[x][1]
            if repl == '-':
                repl = ''
            seq = seq.replace(x, repl)
        pnm = info['pep_nm']
        c.append((-score, mseq, 't' if seq in utils.seen_target else 'd',
            info['charge'], info, info.pop('sumI'), np.median(info.pop('dist')), val['snp_label']))
        c[-1][4]['mzdiff'] = {'Da': s['nm'][info['charge']] - pnm}
        c[-1][4]['mzdiff']['ppm'] = 1e6 * c[-1][4]['mzdiff']['Da'] / pnm
        evalues.append(-1./score if -score else 1e6)
        c = np.array(c, dtype=dtype)
        yield {'spectrum': s, 'candidates': c, 'e-values': evalues}


#!/usr/bin/env python3

import scipy

def find_permutation(n_clusters, real_labels, labels):
    permutation=[]
    for i in range(n_clusters):
        idx = labels == i
        # Choose the most common label among data points in the cluster
        new_label=scipy.stats.mode(real_labels[idx])[0][0]
        permutation.append(new_label)
    return permutation

def plant_clustering():
    return 0.0

def main():
    print(plant_clustering())

if __name__ == "__main__":
    main()


import argparse
import pickle
import numpy as np
from KeywordExtractor import *
from sklearn.feature_extraction import DictVectorizer
from sklearn.pipeline import Pipeline
import datetime
from epitator.annotator import AnnoDoc
from epitator.geoname_annotator import GeonameAnnotator
from epitator.count_annotator import CountAnnotator
from epitator.infection_annotator import InfectionAnnotator
from epitator.date_annotator import DateAnnotator
import disease_label_table
from keyword_annotator import KeywordAnnotator
from epitator.resolved_keyword_annotator import ResolvedKeywordAnnotator
from epitator.structured_incident_annotator import StructuredIncidentAnnotator
from epitator.incident_annotator import IncidentAnnotator

import logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def time_sofar_gen(start_time):
    """
    A generator that returns the time elapsed since the passed in start_time.
    """
    while True:
        yield '[' + str(datetime.datetime.now() - start_time) + ']'

class Diagnoser():

    __version__ = '0.4.4'

    def __init__(
        self, classifier, dict_vectorizer,
        cutoff_ratio=0.65,
        keyword_array=None):
        self.keyword_array = keyword_array
        self.classifier = classifier
        self.geoname_annotator = GeonameAnnotator()
        self.count_annotator = CountAnnotator()
        self.infection_annotator = InfectionAnnotator()
        self.date_annotator = DateAnnotator()
        self.keyword_annotator = KeywordAnnotator()
        self.resolved_keyword_annotator = ResolvedKeywordAnnotator()
        processing_pipeline = []
        processing_pipeline.append(('link', LinkedKeywordAdder(keyword_array)))
        processing_pipeline.append(('limit', LimitCounts(1)))
        self.keyword_processor = Pipeline(processing_pipeline)
        self.dict_vectorizer = dict_vectorizer
        self.keywords = dict_vectorizer.get_feature_names()
        self.keyword_extractor = KeywordExtractor(keyword_array)
        self.cutoff_ratio = cutoff_ratio
    def best_guess(self, X):
        probs = self.classifier.predict_proba(X)[0]
        p_max = max(probs)
        result = {}
        for i,p in enumerate(probs):
            cutoff_ratio = self.cutoff_ratio
            parents = disease_label_table.get_inferred_labels(self.classifier.classes_[i])
            if p >= p_max * self.cutoff_ratio:
                result[i] = max(p, result.get(i, 0))
                for i2, label in enumerate(self.classifier.classes_):
                    if label in parents:
                        result[i2] = max(p, probs[i2], result.get(i2, 0))
        return result.items()
    def diagnose(
        self,
        content,
        diseases_only=False,
        content_date=None,
        use_infection_annotator=False,
        include_incidents=False):
        time_sofar = time_sofar_gen(datetime.datetime.now())
        base_keyword_dict = self.keyword_extractor.transform([content])[0]
        feature_dict = self.keyword_processor.transform([base_keyword_dict])
        X = self.dict_vectorizer.transform(feature_dict)[0]

        logger.info(time_sofar.next() + 'Computed feature vector')
        def diagnosis(i, p):
            scores = self.classifier.coef_[i] * X
            # Scores are normalized so they can be compared across different
            # classifications.
            norm = np.linalg.norm(scores)
            if norm > 0:
                scores /= norm
            scores *= p
            # These might be numpy types. I coerce them to native python
            # types so we can easily serialize the output as json.

            scored_keywords = zip(self.keywords, scores)
            keyword_scores = {}
            for keyword, score in scored_keywords:
                if score > 0 and keyword in base_keyword_dict:
                    keyword_scores[keyword] = float(score)

            return {
                'name': unicode(self.classifier.classes_[i]),
                'probability': float(p),
                'keywords': [{
                        'name': unicode(kwd),
                        'score': float(score),
                    }
                    for kwd, score in scored_keywords
                    if score > 0 and kwd in base_keyword_dict],
                'inferred_keywords': [{
                        'name': unicode(kwd),
                        'score': float(score),
                    }
                    for kwd, score in scored_keywords
                    if score > 0 and kwd not in base_keyword_dict]
            }
        diseases = [diagnosis(i,p) for i,p in self.best_guess(X)]
        if diseases_only:
            return {
                'diseases': diseases
            }
        logger.info(time_sofar.next() + 'Diagnosed diseases')

        anno_doc = AnnoDoc(content, date=content_date)
        anno_doc.add_tier(self.keyword_annotator)
        logger.info('keywords annotated')
        anno_doc.add_tier(self.resolved_keyword_annotator)
        logger.info('resolved keywords annotated')
        anno_doc.add_tier(self.date_annotator)
        logger.info('dates annotated')
        if use_infection_annotator:
            anno_doc.add_tier(self.infection_annotator)
            anno_doc.tiers['counts'] = anno_doc.tiers.pop('infections')
            attribute_remappings = {
                'infection': 'case'
            }
            for span in anno_doc.tiers['counts']:
                span.metadata['attributes'] = [
                    attribute_remappings.get(attribute, attribute)
                    for attribute in span.metadata['attributes']]
        else:
            anno_doc.add_tier(self.count_annotator)
        logger.info('counts annotated')
        anno_doc.add_tier(self.geoname_annotator)
        logger.info('geonames annotated')
        anno_doc.add_tier(StructuredIncidentAnnotator())
        logger.info('structured incidents annotated')
        anno_doc.filter_overlapping_spans(
            tier_names=[ 'dates', 'geonames', 'diseases', 'hosts', 'modes',
                         'pathogens', 'symptoms' ]
        )
        logger.info('filtering overlapping spans done')

        dates = []
        for span in anno_doc.tiers['dates']:
            range_start, range_end = span.datetime_range
            dates.append({
                'type': 'datetime',
                'name': span.text,
                'value': span.text,
                'textOffsets': [
                    [span.start, span.end]
                ],
                'timeRange': {
                    'beginISO': range_start.isoformat().split('T')[0],
                    'begin': {
                        'year': range_start.year,
                        'month': range_start.month,
                        'date': range_start.day
                    },
                    # The date range does not include the end day.
                    'endISO': range_end.isoformat().split('T')[0],
                    'end': {
                        'year': range_end.year,
                        'month': range_end.month,
                        'date': range_end.day
                    },
                }
            })

        geonames_grouped = {}
        for span in anno_doc.tiers['geonames']:
            if not span.geoname['geonameid'] in geonames_grouped:
                geonames_grouped[span.geoname['geonameid']] = {
                    'type': 'location',
                    'name': span.geoname.name,
                    'geoname': span.geoname.to_dict(),
                    'textOffsets': [
                        [span.start, span.end]
                    ]
                }
            else:
                geonames_grouped[
                    span.geoname['geonameid']
                ]['textOffsets'].append(
                    [span.start, span.end]
                )
        logger.info(time_sofar.next() + 'Annotated geonames')

        counts = []
        for span in anno_doc.tiers['counts'].without_overlaps(anno_doc.tiers['structured_data']):
            count_dict = dict(span.metadata)
            count_dict['type'] = 'count'
            count_dict['text'] = span.text
            count_dict['label']= span.label
            count_dict['textOffsets']= [[span.start, span.end]]
            counts.append(count_dict)
            # Include legacy case counts so the diagnositic dashboard
            # doesn't break.
            if 'case' in count_dict['attributes']:
                counts.append({
                    'type': 'caseCount',
                    'text': count_dict['text'],
                    'value': count_dict['count'],
                    'modifiers': count_dict['attributes'],
                    'cumulative': "cumulative" in count_dict['attributes'],
                    'textOffsets': count_dict['textOffsets']
                })
        keyword_types = ['diseases', 'hosts', 'modes', 'pathogens', 'symptoms']
        keyword_groups = {}
        for keyword_type in keyword_types:
            keyword_groups[keyword_type] = {}
            for span in anno_doc.tiers['keyword.' + keyword_type]:
                if span.label not in keyword_groups[keyword_type]:
                    keyword_groups[keyword_type][span.label] = {
                        'type': keyword_type,
                        'value': span.label,
                        'textOffsets': [[span.start, span.end]]
                    }
                else:
                    keyword_groups[keyword_type][span.label]['textOffsets'].append(
                        [span.start, span.end]
                    )
        resolved_keywords = []
        for span in anno_doc.tiers['resolved_keywords'].without_overlaps(anno_doc.tiers['geonames']):
            resolved_keywords.append({
                'type': 'resolvedKeyword',
                'resolutions': span.metadata['resolutions'],
                'text': span.text,
                'textOffsets': [[span.start, span.end]]})
        result = {
            'diagnoserVersion': self.__version__,
            'dateOfDiagnosis': datetime.datetime.now(),
            'diseases': diseases,
            'structuredIncidents': [
                dict(span.metadata, textOffsets=[[span.start, span.end]])
                for span in anno_doc.tiers['structured_incidents']],
            'features': counts +\
                        geonames_grouped.values() +\
                        dates +\
                        keyword_groups['diseases'].values() +\
                        keyword_groups['hosts'].values() +\
                        keyword_groups['modes'].values() +\
                        keyword_groups['pathogens'].values() +\
                        keyword_groups['symptoms'].values() +\
                        resolved_keywords}
        if include_incidents:
            result['incidents'] = []
            anno_doc.add_tier(IncidentAnnotator())
            for incident_span in anno_doc.tiers['incidents']:
                metadata = incident_span.metadata
                incident_data = {
                    'offsets': [span.start, span.end],
                    'type': metadata['type'],
                    'value': metadata['value'],
                    'dateRange': [d.isoformat().split('T')[0] for d in metadata['dateRange']],
                    'locations': metadata['locations'],
                    'species': metadata['species'],
                    'status': metadata.get('status'),
                    'resolvedDisease': metadata.get('resolvedDisease'),
                    'annotations': {
                        'case': [{ 'offsets': [incident_span.start, incident_span.end] }]
                    }
                }
                if 'count_annotation' in metadata:
                    count_annotation = metadata['count_annotation']
                    incident_data['annotations'] = {
                        'case': [{ 'offsets': [count_annotation.start, count_annotation.end] }],
                        'date': [
                            { 'offsets': [anno.start, anno.end] }
                            for anno in metadata['date_territory'].metadata
                        ],
                        'location': [
                            { 'offsets': [anno.start, anno.end] }
                            for anno in metadata['geoname_territory'].metadata
                        ],
                        'disease': [
                            { 'offsets': [anno.start, anno.end] }
                            for anno in metadata['disease_territory'].metadata
                        ]
                    }
                result['incidents'].append(incident_data)
        return result

if __name__ == '__main__':
    import Diagnoser
    parser = argparse.ArgumentParser()
    parser.add_argument('content', metavar='content', type=str, help='Text to diagnose')
    args = parser.parse_args()
    content = args.content
    with open('diagnoser.p', 'rb') as f:
        my_diagnoser = pickle.load(f)
        print my_diagnoser.diagnose(content)