import os
import argparse
import datetime
import torch
from bert_serving.client import BertClient
from torch.utils.data import (DataLoader, RandomSampler, SequentialSampler,
                              TensorDataset)
from torch.utils.data.distributed import DistributedSampler
import logging
from tqdm import tqdm, trange
from torch.autograd import Variable
import torch.nn as nn
import torch.nn.functional as F
import random
import numpy as np
import csv
import json
import data
import model
import train

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='RCNN text classificer')
    # learning
    parser.add_argument('-lr', type=float, default=1e-3, help='initial learning rate [default: 0.001]')
#     parser.add_argument('-per-gpu-train-batch-size', type=int, default=2, help='per gpu batch size for training [default: 2]')
#     parser.add_argument('-per-gpu-eval-batch-size', type=int, default=32, help='per gpu batch size for eval [default: 32]')
    parser.add_argument('-batch-size', type=int, default=32, help='batch size for train [default: 32]')
    parser.add_argument('-save', type=int, default=1, help='whether to save model')
    parser.add_argument('-model', type=str, default=None, required=True, help='dir to store model')
    parser.add_argument('-predict-file', type=str, default=None, required=True, help='dir to store predict')
    # data
    parser.add_argument('-data-dir', type=str, default='../data', help='dataset dir')
    parser.add_argument('-train-file', type=str, default=None, required=True, help='train file')
    parser.add_argument('-dev-file', type=str, default=None, required=True, help='dev file')
    parser.add_argument('-test-file', type=str, default=None, required=True, help='test file')
    parser.add_argument('-eval-result', type=str, default=None, required=True, help='eval result file')
    # model
    parser.add_argument('-dropout', type=float, default=0.5, help='the probability for dropout [default: 0.5]')
    parser.add_argument('-input-size', type=int, default=768, help='input dimensiton')
    parser.add_argument('-hidden-size', type=int, default=100, help='rnn hidden layer size')
    parser.add_argument('-seq-len', type=int, default=400, help='input seq len')
    parser.add_argument('-hidden_layers', type=int, default=1, help='rnn hidden layers')
    parser.add_argument('-directions', type=int, default=2, help='bidirection rnn')
    parser.add_argument('-concat', type=int, default=None, required=True, help='title&content concat type')
    parser.add_argument('-linear-size', type=int, default=100, help='linear1 dimensiton')
    parser.add_argument('-train-steps', type=int, default=100, help='train steps, one step for one batch')
    parser.add_argument('-model-type', type=str, default='rnn', help='recurrent network model type')
    # device
    parser.add_argument('-server-ip', type=str, default=None, required=True, help='device ip for bert-as-service server')
    parser.add_argument('-do-predict', type=int, default=0, help='predict the sentence given')
    parser.add_argument('-do-train', type=int, default=1, help='train or test')
    parser.add_argument('-load-model', type=str, default=None, required=True, help='predict loading model dir')
    parser.add_argument('-port', type=int, default=5555, help='bert-as-service port')
    parser.add_argument('-port-out', type=int, default=5556, help='bert-as-service port')
    args = parser.parse_args()

    # connect to bert-serving-server
    bc = BertClient(ip=args.server_ip,port=args.port, port_out=args.port_out, check_version=False, check_length=False)
    mydata = data.MyTaskProcessor()
    logging.basicConfig(format = '%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
                    datefmt = '%m/%d/%Y %H:%M:%S',
                    level = logging.INFO)
    logger = logging.getLogger(__name__)
    # set-up cuda,gpu
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # update args
    args.class_num = len(mydata.get_labels()) # 3

    # define text rcnn model
    print('lr: '+str(args.lr))
    rcnn = model.RCNN_Text(args.input_size, args.hidden_size, args.class_num, args.input_size+2*args.hidden_size, args.linear_size, args.dropout, args.model_type).to(device) #need to modify

    # training
    if args.do_train:

        # open eval_result.txt
        with open(args.eval_result, 'w') as f:
            now = datetime.datetime.now()
            f.write('['+str(now)+']\n')
            f.write('*'*80+'\n')

        # get eval data&label
        dev, dev_label = mydata.get_dev_examples(args.data_dir, args.batch_size, bc, args.dev_file, args.concat)
        train_, train_label = mydata.get_train_examples(args.data_dir, args.batch_size, bc, args.train_file, args.concat)
        print('='*80)
        logger.info("*** Start Training ***")
        logger.info("  Num examples = %d", (len(train_)-1)*args.batch_size+len(train_[-1]))
        logger.info("  Batch size = %d", len(train_))
        logger.info("  Num steps = %d", args.train_steps)

        # set tqdm bar
        num_train_optimization_steps = args.train_steps
        bar = tqdm(range(num_train_optimization_steps),total=num_train_optimization_steps)

        # loss_func define
        loss_func = nn.CrossEntropyLoss(weight=torch.from_numpy(np.array([1.19462648, 0.25, 0.310986013])).float(), size_average=True).to(device)
        optimizer = torch.optim.Adam(rcnn.parameters(), lr=args.lr)
        max_f1 = 0

        # train progress
        for i, step in enumerate(bar):
            # get batch data
            batch_train = train_[int(i%len(train_))]
            batch_label = train_label[int(i%len(train_label))]
            # train
            train_loss = train.train(rcnn, train=batch_train, train_label=batch_label, loss_func=loss_func, optimizer=optimizer, device=device, eval_result=args.eval_result, hidden_size=args.hidden_size, seq_len=args.seq_len, input_size=args.input_size, linear_size=args.linear_size, step=i, model_type=args.model_type)
            bar.set_description("loss {}".format(train_loss))

            # eval
            macro_f1 = train.eval(rcnn, dev=dev, dev_label=dev_label, batch_size=args.batch_size, device=device, num_label=args.class_num, eval_result=args.eval_result, hidden_size=args.hidden_size, seq_len=args.seq_len, input_size=args.input_size, linear_size=args.linear_size, step=i, model_type=args.model_type)

            if macro_f1 >= max_f1:
                print("Best F1", macro_f1)
                print("Saving Model......")
                max_f1=macro_f1
                # Save a trained model
                if args.save == 1:
                    torch.save(rcnn.state_dict(), args.model) #保存网络参数
                    print("="*80)
            else:
                print("="*80)
        logger.info("  Max Macro_f1 = %f", max_f1)
        print('='*80)
        with open(args.eval_result, 'a') as f:
            now = datetime.datetime.now()
            f.write('['+str(now)+'] max_macro_f1:' + str(max_f1))

    # predicting
    if args.do_predict:
        rcnn.load_state_dict(torch.load(args.load_model))
        print('model load done')
        test = mydata.get_test_examples(args.data_dir, args.batch_size, bc, args.test_file, args.concat)
        print('test data done')
        train.predict(rcnn, test=test, device=device, file=args.predict_file, hidden_size=args.hidden_size, seq_len=args.seq_len, input_size=args.input_size, linear_size=args.linear_size, model_type=args.model_type)
        print('predict saved')
        print('='*80)
    bc.close()


# Standard library
import os

# Third-party
import numpy as np
from numpy.lib.recfunctions import append_fields
from scipy.integrate import quad
from scipy.interpolate import interp1d
from astropy.table import Table
from astropy import units as u

# Project
from ._read_mist_models import IsoCmdReader, IsoReader
from .imf import imf_dict, IMFIntegrator
from .. import MIST_PATH
from ..log import logger
from ..filters import phot_system_list, get_filter_names
from ..filters import load_zero_point_converter
from ..util import check_units, fetch_mist_grid_if_needed
phot_str_helper = {p.lower():p for p in phot_system_list}

__all__ = ['fetch_mist_iso_cmd', 'Isochrone', 'MISTIsochrone']

class Isochrone(object):
    """
    Class for storing generic isochrones. It also has methods for calculating
    IMF-weighted properties of a simple stellar population with the given
    age and metallicity.

    Parameters
    ----------
    mini : `~numpy.ndarray`
        Initial stellar masses. The masses must be monotonically increasing.
    mact : `~numpy.ndarray`
        Actual stellar masses after mass loss.
    mag_table : `~astropy.table.Table`, dict, or structured `~numpy.ndarray`
        The stellar magnitudes.
    eep : `~numpy.ndarray`, optional
        The Primary Equivalent Evolutionary Points. Needed to identify phases
        of stellar evolution.
    log_L : `~numpy.ndarray`, optional
        Stellar luminosities.
    log_Teff : `~numpy.ndarray`, optional
        Stellar effective temperatures.
    """

    def __init__(self, mini, mact, mags, eep=None, log_L=None,
                 log_Teff=None):
        self.mini = np.asarray(mini)
        self.mact = np.asarray(mact)
        if (np.diff(self.mini) < 0).sum() > 0:
            raise Exception('Initial masses must be monotonically increasing.')
        self.eep = None if eep is None else np.asarray(eep)
        self.log_L = None if log_L is None else np.asarray(log_L)
        self.log_Teff = None if log_Teff is None else np.asarray(log_Teff)
        if type(mags) == dict or type(mags) == np.ndarray:
            self.mag_table = Table(mags)
        elif type(mags) == Table:
            self.mag_table = mags
        else:
            raise Exception(f'{type(mags)} is not a valid type for mags')

    @property
    def m_min(self):
        """The minimum mass of the isochrone."""
        return self.mini.min()

    @property
    def m_max(self):
        """The maximum mass of the isochrone."""
        return self.mini.max()

    @property
    def filters(self):
        """List of filters in the given photometric system(s)."""
        return self.mag_table.colnames

    @staticmethod
    def from_parsec(file_name, log_age=None, zini=None, num_filters=None):
        """
        Read isochrone generated from the `PARSEC website
        <http://stev.oapd.inaf.it/cgi-bin/cmd>`_. If more than one age and/or
        metallicity is included in the file, you must provide the log_age
        and/or zini parameters.

        Parameters
        ----------
        file_name : str
            Isochrone file name.
        log_age : float, optional
            Log of age in years. You must provided this parameter if there is
            more than one age in the isochrone file. Note this function does
            not interpolate ages, so it must be included in the file.
        zini : float, optional
            Initial metal fraction. You must provided this parameter if there
            is more than one metallicity in the isochrone file. Note this
            function does not interpolate metallicity, so it must be included
            in the file.
        num_filters : int, optional
            Number of filters included in the isochrone file. If None, will
            assume the last non-filter parameter is `mbolmag`.

        Returns
        -------
        iso : `artpop.stars.Isochrone`
            PARSEC Isochrone object.
        """
        if os.path.isfile(file_name):
            file = open(file_name, 'r')
            lines = file.readlines()
            file.close()
            for l in lines:
                if 'Zini' in l.split()[1]:
                    names = l.split()[1:]
                    break
            data = np.loadtxt(file_name)
            parsec = Table(data, names=names)
        else:
            raise Exception(f'{file_name} does not exist.')
        isochrone_full = parsec.copy()
        if log_age is not None:
            age_cut = np.abs(parsec['logAge'] - log_age) < 1e-5
            if age_cut.sum() < 1:
                raise Exception(f'log_age = {log_age} not found.')
            parsec = parsec[age_cut]
        if zini is not None:
            zini_cut = np.abs(parsec['Zini'] - zini) < 1e-8
            if zini_cut.sum() < 1:
                raise Exception(f'Zini = {zini} not found.')
            parsec = parsec[zini_cut]
        if num_filters is None:
            filt_idx = np.argwhere(np.array(names) == 'mbolmag')[0][0] + 1
        else:
            filt_idx = len(names) - num_filters
        iso = Isochrone(mini=parsec['Mini'],
                        mact=parsec['Mass'],
                        mags=parsec[names[filt_idx:]],
                        log_L=parsec['logL'],
                        log_Teff=parsec['logTe'])
        iso.isochrone_full = isochrone_full
        return iso

    def interpolate(self, y_name, x_interp, x_name='mini',
                    slice_interp=np.s_[:], **kwargs):
        """
        Interpolate isochrone.

        Parameters
        ----------
        y_name : str
            Parameter name of interpolated variable.
        x_interp : `~numpy.ndarray`
            New x values to interpolate over.
        x_name : str
            Parameter name of x values. Usually want this to be initial mass.

        Returns
        -------
        y_interp : `~numpy.ndarray`
            Interpolated y values.
        """
        x = self.mag_table[x_name] if x_name in self.filters\
                                   else getattr(self, x_name)
        y = self.mag_table[y_name] if y_name in self.filters\
                                   else getattr(self, y_name)
        x = x[slice_interp]
        y = y[slice_interp]
        y_interp = interp1d(x, y, **kwargs)(x_interp)
        return y_interp

    def mag_to_mass(self, mag, bandpass):
        """
        Interpolate isochrone to calculate initial stellar masses that have
        the given magnitude. Since more than one mass can have the same
        magnitude (from stars being in different phases of stellar evolution),
        we step through each isochrone and check if the given magnitude falls
        between two bins, in which case we interpolate the associated mass.

        Parameters
        ----------
        mag : float
            Magnitude to be converted in to stellar mass(es).
        bandpass : str
            Photometric bandpass of the `mag`.

        Returns
        -------
        mass_interp : `~numpy.ndarray`
            Interpolated mass(es) associated with `mag`.
        """
        y = self.mini
        x = self.mag_table[bandpass]

        if mag < x.min() or mag > x.max():
            raise Exception(f'mag = {mag} is outside the isochrone range.')

        y_interp = []
        # Loop over isochrone and check if desired magnitude is in between two
        # bins. If it is, interpolate the mass. Looping like this is necessary
        # because different masses can have the same magnitude due to
        # being at different phases of stellar evolution.
        for i in range(len(x)):
            if x[i] == mag:
                y_interp.append(y[i])
                continue
            if i < len(x) - 1:
                _x = [x[i], x[i + 1]]
                _y = [y[i], y[i + 1]]
                dx = x[i + 1] - x[i]
                if dx > 0:
                    if (mag <= x[i + 1]) and (mag > x[i]):
                        y_interp.append(interp1d(_x, _y)([mag])[0])
                else:
                    if (mag >= x[i + 1]) and (mag < x[i]):
                        y_interp.append(interp1d(_x, _y)([mag])[0])
        mass_interp = np.array(y_interp)

        return mass_interp

    def calculate_mag_limit(self, imf, bandpass, frac_mass_sampled=None,
                            frac_num_sampled=None, distance=10*u.pc):
        """
        Calculate the limiting faint magnitude to sample a given fraction
        of mass or number of a stellar population. This is used for when you
        only what to sample stars that are brighter than a magnitude limit.
        You must specify either `frac_mass_sampled` or `frac_num_sampled`.

        .. note::
            You must give `frac_mass_sampled` *or* `frac_num_sampled`.

        Parameters
        ----------
        imf : str
            Name for the IMF.
        bandpass : str
            Observation bandpass of the limiting magnitude.
        frac_mass_sampled: float, optional
            Fraction of mass that will be sampled if only stars that are
            brighter than the magnitude limit are sampled.
        frac_num_sampled: float, optional
            Fraction of stars by number that will be sampled if only stars that
            are brighter than the magnitude limit are sampled.
        distance : float or `~astropy.units.Quantity`, optional
            Distance to source. If float is given, the units are assumed
            to be `~astropy.units.Mpc`. Default distance is 10
            `~astropy.units.pc` (i.e., the mags are in absolute units).

        Returns
        -------
        mag_limit : float
            The desired limiting magnitude.
        mass_limit : float
            Mass of stars (in solar masses) that have the limiting magnitude.
        """
        mfint = IMFIntegrator(imf, self.m_min, self.m_max)
        m_vals = np.logspace(np.log10(self.m_min), np.log10(self.m_max), 500)
        err_msg = 'You must give a fraction such that 0 < fraction <= 1.'
        if frac_mass_sampled is not None:
            assert frac_mass_sampled > 0 and frac_mass_sampled <= 1, err_msg
            frac_interp = frac_mass_sampled
            fracs = [mfint.m_integrate(_m, self.m_max, True) for _m in m_vals]
        elif frac_num_sampled is not None:
            assert frac_num_sampled > 0 and frac_num_sampled <= 1, err_msg
            frac_interp = frac_num_sampled
            fracs = [mfint.integrate(_m, self.m_max, True) for _m in m_vals]
        else:
            msg = 'You must give either frac_mass_sampled or frac_num_sampled.'
            raise Exception(msg)
        d = check_units(distance, 'Mpc')
        mass_limit = interp1d(fracs, m_vals)([frac_interp])[0]
        mag_limit = self.interpolate(bandpass, [mass_limit])[0]
        mag_limit = mag_limit + 5 * np.log10(d.to('pc').value) - 5
        return mag_limit, mass_limit

    def nearest_mini(self, m):
        """
        Find the nearest mass to `m` in the initial mass array.

        Parameters
        ----------
        m : float:
            The mass for which we want the nearest initial mass.

        Returns
        -------
        m_nearest : float
            Value of the nearest mass.
        arg_nearest : int
            Argument of the nearest mass.
        """
        arg_nearest = np.abs(self.mini - m).argmin()
        m_nearest = self.mini[arg_nearest]
        return m_nearest, arg_nearest

    def imf_weights(self, imf, m_min_norm=None, m_max_norm=None,
                    norm_type='mass', **kwargs):
        """
        Calculate IMF weights.

        Parameters
        ----------
        imf : str or callable function that takes mass as its argument
            Stellar Initial mass function.
        m_min_norm : None or float, optional
            Minimum mass for the normalization. Must be less than or equal to
            the mini mass of isochrone, which will be used if None is given.
        m_max_norm : float, optional
            Maximum mass for normalization.
        norm_type : str, optional
            Type of IMF normalization (mass or number).

        Returns
        -------
        wght : `~numpy.ndarray`
            IMF weights calculated such that an integral over mass is simply
            given by SUM(m * wght).
        """
        m_min = m_min_norm if m_min_norm else self.mini.min()
        m_max = m_max_norm if m_max_norm else self.mini.max()
        if m_min > self.m_min:
            raise Exception('Minimum mass must be <= isochrone min mass.')
        if callable(imf):
            imf_func = imf
        else:
            imf_func = lambda m: imf_dict[imf](m, norm_type=None)
        m_imf_func = lambda m: m * imf_func(m)
        norm_func = dict(mass=m_imf_func, number=imf_func)[norm_type]
        norm = quad(norm_func, m_min, m_max, **kwargs)[0]

        wght = []
        mini = self.mini
        # Assume mass is constant in each bin and integrate.
        # This means an integral over mass is simply SUM(m * wght)
        # This trick was stolen from Charlie Conroy's FSPS code :)
        for i in range(len(mini)):
            if i == 0:
                m1 = m_min
            else:
                m1 = mini[i] - 0.5 * (mini[i] - mini[i-1])
            if i == len(mini) - 1:
                m2 = mini[i]
            else:
                m2 = mini[i] + 0.5 * (mini[i+1] - mini[i])
            if m2 < m1:
                raise Exception('Masses must be monotonically increasing.')
            wght.append(quad(imf_func, m1, m2, **kwargs)[0])
        wght = np.array(wght) / norm
        return wght

    def ssp_color(self, blue, red, imf='kroupa', **kwargs):
        """
        Calculate IMF-weighted integrated color.

        Parameters
        ----------
        blue : str
            Blue bandpass.
        red : str
            Red bandpass.
        imf : str, optional
            IMF name.

        Returns
        -------
        color : float
            Integrated color of stellar population.
        """
        wght = self.imf_weights(imf, **kwargs)
        lum_blue = np.sum(wght * 10**(-0.4 * self.mag_table[blue]))
        lum_red = np.sum(wght * 10**(-0.4 * self.mag_table[red]))
        color = -2.5 * np.log10(lum_blue / lum_red)
        return color

    def ssp_sbf_mag(self, bandpass, imf='kroupa', **kwargs):
        """
        Calculate IMF-weighted SBF magnitude.

        Parameters
        ----------
        bandpass : str
            Bandpass to of SBF mag.
        imf : str, optional
            IMF name.

        Returns
        -------
        sbf : float
            SBF magnitude of stellar population.
        """
        wght = self.imf_weights(imf, **kwargs)
        lumlum = np.sum(wght * 10**(-0.8 * self.mag_table[bandpass]))
        lum = np.sum(wght * 10**(-0.4 * self.mag_table[bandpass]))
        sbf = -2.5 * np.log10(lumlum / lum)
        return sbf

    def ssp_mag(self, bandpass, imf='kroupa', norm_type='mass', **kwargs):
        """
        Calculate IMF-weighted magnitude.

        Parameters
        ----------
        bandpass : str
            Bandpass to of mean mag.
        imf : str, optional
            IMF name.
        norm_type : str, optional
            Normalization type (mass or number)

        Returns
        -------
        m : float
            Integrated magnitude of stellar population.
        """
        wght = self.imf_weights(imf, norm_type=norm_type, **kwargs)
        lum = np.sum(wght * 10**(-0.4 * self.mag_table[bandpass]))
        m = -2.5 * np.log10(lum)
        return m

    def ssp_mean_star_mass(self, imf):
        """
        Calculate IMF-weighted mean stellar mass, where the IMF is normalized
        by number to one star formed. The mean stellar mass is equivalent to
        dividing a large of samples (distributed according to the IMF) by the
        number of stars.

        Parameters
        ----------
        imf : str:
            Initial mass function name.

        Returns
        -------
        mean_mass : float
            The mean stellar mass.
        """
        w = self.imf_weights(imf, norm_type='number',
                             m_max_norm=self.mini.max())
        mean_mass = np.sum(self.mact * w)
        return mean_mass

    def ssp_surviving_mass(self, imf, m_min=None, m_max=120,
                           add_remnants=True, mlim_bh=40.0, mlim_ns=8.5):
        """
        Calculate IMF-weighted mass normalized such that 1 solar mass is
        formed by the age of the population.

        The initial-mass-dependent remnant formulae are taken
        from `Renzini & Ciotti 1993
        <https://ui.adsabs.harvard.edu/abs/1993ApJ...416L..49R/abstract>`_.

        Parameters
        ----------
        imf : str
            IMF name.
        add_remnants : bool
            If True, add stellar remnants to surviving mass.

        Returns
        -------
        mass : float
            Total surviving mass.
        """
        m_min = m_min if m_min else self.mini.min()
        m_max = m_max if m_max else self.mini.max()
        wght = self.imf_weights(imf, m_min_norm=m_min, m_max_norm=m_max)
        mass = np.sum(self.mact * wght)

        if add_remnants:
            imf_func = lambda m: imf_dict[imf](m, norm_type=None)
            m_imf_func = lambda m: m * imf_func(m)
            norm = quad(m_imf_func, m_min, m_max)[0]

            # BH remnants
            m_low = max(mlim_bh, self.m_max)
            mass = mass + 0.5 * quad(m_imf_func, m_low, m_max)[0] / norm

            # NS remnants
            if self.m_max <= mlim_bh:
                m_low = max(mlim_ns, self.m_max)
                mass = mass + 1.4 * quad(imf_func, m_low, mlim_bh)[0] / norm

            # WD remnants
            if self.m_max < mlim_ns:
                mmax = self.m_max
                mass = mass + 0.48 * quad(imf_func, mmax, mlim_ns)[0] / norm
                mass = mass + 0.077 * quad(m_imf_func, mmax, mlim_ns)[0] / norm

        return mass

def fetch_mist_iso_cmd(log_age, feh, phot_system, mist_path=MIST_PATH,
                       v_over_vcrit=0.4):
    """
    Fetch MIST isochrone grid.

    Parameters
    ----------
    log_age : float
        Logarithm base 10 of the simple stellar population age in years.
    feh : float
        Metallicity [Fe/H] of the simple stellar population.
    phot_system : str
        Name of the photometric system.
    mist_path : str, optional
        Path to MIST isochrone grids. Use this if you want to use a different
        path from the `MIST_PATH` environment variable.
    v_over_vcrit : float, optional
        Rotation rate divided by the critical surface linear velocity. Current
        options are 0.4 (default) and 0.0.

    Returns
    -------
    iso_cmd : `~numpy.ndarray`
        Structured ``numpy`` array with isochrones and stellar magnitudes.
    """

    # fetch the mist grid if necessary
    fetch_mist_grid_if_needed(phot_system, v_over_vcrit, mist_path)

    v = f'{v_over_vcrit:.1f}'
    ver = 'v1.2'
    p = phot_str_helper[phot_system.lower()]
    path = os.path.join(mist_path, 'MIST_' + ver + f'_vvcrit{v}_' + p)
    sign = 'm' if feh < 0 else 'p'
    fn = f'MIST_{ver}_feh_{sign}{abs(feh):.2f}_afe_p0.0_vvcrit{v}_{p}.iso.cmd'
    fn = os.path.join(path, fn)
    iso_cmd = IsoCmdReader(fn, verbose=False)
    iso_cmd = iso_cmd.isocmds[iso_cmd.age_index(log_age)]
    return iso_cmd

class MISTIsochrone(Isochrone):
    """
    Class for fetching and storing MIST isochrones. It also has several methods
    for calculating IMF-weighted photometric properties of a stellar population
    with the given age an metallicity.

    .. note::
        Currently, the models are interpolated in metallicity but not in age.
        Ages are therefore limited to the age grid of the MIST models. The
        [Fe/H] and log(Age/yr) grids are stored as the private class attributes
        ``_feh_grid`` and ``_log_age_grid``.

    Parameters
    ----------
    log_age : float
        Logarithm base 10 of the simple stellar population age in years.
    feh : float
        Metallicity [Fe/H] of the simple stellar population.
    phot_system : str or list-like
        Name of the photometric system(s).
    mist_path : str, optional
        Path to MIST isochrone grids. Use this if you want to use a different
        path from the ``MIST_PATH`` environment variable.
    ab_or_vega : str, optional
        Magnitudes will be in the AB (default) or Vega magnitude system.
    v_over_vcrit : float, optional
        Rotation rate divided by the critical surface linear velocity. Current
        options are 0.4 (default) and 0.0.
    """

    # the age grid
    _log_age_grid = np.arange(5.0, 10.3, 0.05)
    _log_age_min = _log_age_grid.min()
    _log_age_max = _log_age_grid.max()

    # the [Fe/H] metallicity grid
    # mist has feh <= -4, but using <=-3 due to interpolation issues
    _feh_grid = np.concatenate([np.arange(-3.0, -2., 0.5),
                                np.arange(-2.0, 0.75, 0.25)])
    _feh_min = _feh_grid.min()
    _feh_max = _feh_grid.max()

    def __init__(self, log_age, feh, phot_system, mist_path=MIST_PATH,
                 ab_or_vega='ab', v_over_vcrit=0.4):

        # verify age are metallicity are within model grids
        if log_age < self._log_age_min or log_age > self._log_age_max:
            raise Exception(f'log_age = {log_age} not in range of age grid')
        if feh < self._feh_min or feh > self._feh_max:
            raise Exception(f'feh = {feh} not in range of feh grid')

        self.feh = feh
        self.mist_path = mist_path
        self.phot_system = phot_system
        self.ab_or_vega = ab_or_vega
        self.v_over_vcrit = v_over_vcrit

        # use nearest age (currently not interpolating on age)
        age_diff = np.abs(self._log_age_grid - log_age)
        self.log_age = self._log_age_grid[age_diff.argmin()]
        if age_diff.min() > 1e-6:
            logger.debug('Using nearest log_age = {:.2f}'.format(self.log_age))

        # store phot_system as list to allow multiple photometric systems
        if type(phot_system) == str:
            phot_system = [phot_system]

        # fetch first isochrone grid, interpolating on [Fe/H] if necessary
        self._iso_full = self._fetch_iso(phot_system[0])

        # iterate over photometric systems and fetch remaining isochrones
        filter_dict = get_filter_names()
        filters = filter_dict[phot_system[0]].copy()
        for p in phot_system[1:]:
            filt = filter_dict[p].copy()
            filters.extend(filt)
            _iso = self._fetch_iso(p)
            mags = [_iso[f].data for f in filt]
            self._iso_full = append_fields(self._iso_full, filt, mags)

        # covert magnitudes to ab or vega as necessary
        self.zpt_convert = load_zero_point_converter()
        for filt in filters:
            converter = getattr(self.zpt_convert, f'to_{ab_or_vega.lower()}')
            try:
                m_convert = converter(filt)
            except AttributeError:
                m_convert = 0.0
                logger.warning(f'No AB / Vega conversion found for {filt}.')
            self._iso_full[filt] = self._iso_full[filt] + m_convert

        super(MISTIsochrone, self).__init__(
            mini = self._iso_full['initial_mass'],
            mact = self._iso_full['star_mass'],
            mags = Table(self._iso_full[filters]),
            eep = self._iso_full['EEP'],
            log_L = self._iso_full['log_L'],
            log_Teff = self._iso_full['log_Teff'],
        )

    @property
    def isochrone_full(self):
        """MIST entire isochrone in a structured `~numpy.ndarray`."""
        return self._iso_full

    @staticmethod
    def from_parsec(fn, **kwargs):
        msg = 'PARSEC isochrones do not with MISTIsochrone.'
        raise Exception(msg + ' Use artpop.Isochrone instead.')

    def _fetch_iso(self, phot_system):
        """Fetch MIST isochrone grid, interpolating on [Fe/H] if necessary."""
        if self.feh in self._feh_grid:
            args = [self.log_age, self.feh, phot_system, self.mist_path,
                    self.v_over_vcrit]
            iso = fetch_mist_iso_cmd(*args)
        else:
            iso = self._interp_on_feh(phot_system)
        return iso

    def _interp_on_feh(self, phot_system):
        """Interpolate isochrones between two [Fe/H] grid points."""
        i_feh = self._feh_grid.searchsorted(self.feh)
        feh_lo, feh_hi = self._feh_grid[i_feh - 1: i_feh + 1]

        logger.debug('Interpolating to [Fe/H] = {:.2f} '\
                     'using [Fe/H] = {} and {}'.\
                     format(self.feh, feh_lo, feh_hi))

        mist_0 = fetch_mist_iso_cmd(
            self.log_age, feh_lo, phot_system, self.mist_path)
        mist_1 = fetch_mist_iso_cmd(
            self.log_age, feh_hi, phot_system, self.mist_path)

        y0, y1 = np.array(mist_0.tolist()), np.array(mist_1.tolist())

        x = self.feh
        x0, x1 = feh_lo, feh_hi
        weight = (x - x0) / (x1 - x0)

        len_0, len_1 = len(y0), len(y1)

        # if necessary, extrapolate using trend of the longer array
        if (len_0 < len_1):
            delta = y1[len_0:] - y1[len_0 - 1]
            y0 = np.append(y0, y0[-1] + delta, axis=0)
        elif (len_0 > len_1):
            delta = y0[len_1:] - y0[len_1 - 1]
            y1 = np.append(y1, y1[-1] + delta, axis=0)

        y = y0 * (1 - weight) + y1 * weight
        iso = np.core.records.fromarrays(y.transpose(), dtype=mist_0.dtype)

        return iso

import numpy,from numpy,from scipy,from astropy


# Adapted from repo botwizer by DevGltich
# https://github.com/DevGlitch/botwizer

# Resources used:
# https://github.com/AlexeyAB/darknet

import cv2
import time
import numpy as np

def vid_object_detection_txt(vid_path, config_path, weights_path, labels_path):
    """Running YOLO on a video to detect objects
    :param vid_path: path of video to analyse
    :param config_path: path of the .cfg file
    :param weights_path: path of the .weights file
    :param labels_path: path of the .names file
    :return: object(s) detected
    :rtype: list
    """

    # Reads and load model stored in Darknet model files
    net = cv2.dnn.readNetFromDarknet(config_path, weights_path)

    # Object Labels
    obj_names = open(labels_path)
    obj_labels = obj_names.read().strip().split("\n")

    # initialize a list of colors to represent each possible class label
    # np.random.seed(14)
    colors = np.random.randint(0, 255, size=(len(obj_labels), 3), dtype="uint8")

    # Reads video from provided path
    # vid_file = os.path.join(vid_path)
    vid_cap = cv2.VideoCapture(vid_path)

    # For writing output video
    _, image = vid_cap.read()
    img_row, img_col = image.shape[:2]
    fourcc = cv2.VideoWriter_fourcc(*"XVID")
    out = cv2.VideoWriter("output.mp4", fourcc, 20.0, (img_col, img_row))

    # # List objects where it stores the obj_names labels detected in the image
    # objects = []

    while vid_cap.isOpened():

        _, image = vid_cap.read()
        img_row, img_col = image.shape[:2]

        # Creating a 4-dimensional blob from image
        # SwapRB to True increase classification accuracy
        blob = cv2.dnn.blobFromImage(
            image, 1 / 255.0, (416, 416), swapRB=True, crop=False
        )
        net.setInput(blob)

        # Putting blob as the input of the network
        net.setInput(blob)

        # Getting each layer name
        layer_name = net.getLayerNames()
        layer_name = [layer_name[i[0] - 1] for i in net.getUnconnectedOutLayers()]

        # # Timing network output
        # time_start = time.perf_counter()
        # print("Starting YOLO analysis...")
        outputs = net.forward(layer_name)
        # time_stop = time.perf_counter() - time_start
        # print(f"YOLO ran for: {time_stop:.2f}s")

        grid, probabilities, labels = [], [], []

        # Find each single output
        # This for loop is based on information from darknet's code and opencv
        for output in outputs:

            # Find each single detection in output
            for detection in output:

                # Get probability score and label of the detection
                score = detection[5:]
                label = np.argmax(score)
                prob = score[label]

                # Selecting only detections that are superior to 50% probability
                # Anything below 50% is ignored as probability is too low
                # You can increase this to higher or lower probability if needed
                if prob > 0.5:

                    # Working on each bounding box of the grid created by YOLO
                    grid_box = detection[:4] * np.array(
                        [img_col, img_row, img_col, img_row]
                    )
                    (X, Y, width, height) = grid_box.astype("int")
                    x = X - (width / 2)
                    y = Y - (height / 2)

                    # Appending to the lists
                    probabilities.append(float(prob))
                    labels.append(label)
                    grid.append([int(x), int(y), int(width), int(height)])

        # Performs Non Maximum Suppression given boxes and corresponding scores.
        # This filters the boxes in the image grid.
        # It keeps only the ones with the highest probability
        NMS = cv2.dnn.NMSBoxes(grid, probabilities, 0.6, 0.6)

        # If at least one object has been detected
        if len(NMS) > 0:

            # List objects where it stores the obj_names labels detected in the image
            objects = []

            for i in NMS.flatten():

                objects += [f"{obj_labels[labels[i]]}"]

                # print("The objects are:\n", objects)

                yield objects

    # Close names file
    obj_names.close()

    # Close video file
    vid_cap.release()

    # return objects


import subprocess
import random
import sys
import os
from datetime import datetime
import argparse
import atexit
import numpy as np
from scipy.io import wavfile
import librosa

# input params
parser = argparse.ArgumentParser(description="Removes silence from a video")
parser.add_argument("input_file", type=str, help="the video file you want modified")
parser.add_argument("-db", type=float, default=40, help="db threshold of the parts to remove")
parser.add_argument("-out_format", type=str, default="mp4", help="format of output file")
parser.add_argument("-w", action="store_true", help="directly watch video with vlc")
parser.add_argument("-vlc_path", type=str, default="/Applications/VLC.app/Contents/MacOS/VLC", help="path to vlc")

args = parser.parse_args()

salt = f"{random.randint(0, 10**5):05d}"
audio_file = f"JC_AUDIO_{salt}.wav"
script_file = f"JC_SCRIPT_{salt}.txt"

if args.w:
    args.out_format = "ts"

def run_ffmpeg(cmd):
    p = subprocess.run(["ffmpeg", *cmd])
    if p.returncode:
        print(f"ffmpeg terminated with code {p.returncode}.", file=sys.stderr)
        sys.exit(1)

def cleanup():
    try:
        os.remove(audio_file)
        os.remove(script_file)
    except:
        pass

atexit.register(cleanup)

start_time = datetime.now()

# extract audio
run_ffmpeg(["-i", args.input_file, "-acodec", "pcm_s16le", "-ac", "1", "-ar", "16000", audio_file])
sample_rate, audio_data = wavfile.read(audio_file)

# select sounded parts
sounded_parts = librosa.effects.split(
    audio_data.astype(np.float32),
    top_db=args.db,
    frame_length=2048,
    hop_length=512,
)

# generate script
with open(script_file, "w") as f:
    for i, (a, b) in enumerate(sounded_parts):
        start = round(a / sample_rate, 3)
        end = round(b / sample_rate, 3)
        f.write(f"[0:v]trim=start={start}:end={end},setpts=PTS-STARTPTS[v{i}];\n")
        f.write(f"[0:a]atrim=start={start}:end={end},asetpts=PTS-STARTPTS[a{i}];\n")

    n = len(sounded_parts)
    streams_v = "".join(map(lambda x: f"[v{x}]", range(n)))
    streams_a = "".join(map(lambda x: f"[a{x}]", range(n)))
    f.write(f"{streams_v}concat=n={n}[vout];\n")
    f.write(f"{streams_a}concat=n={n}:v=0:a=1[aout]")

# render video
frac = int((sounded_parts[:, 1] - sounded_parts[:, 0]).sum() / len(audio_data) * 100)
out_file = args.input_file.rsplit(".", 1)[0] + f" ({frac}%).{args.out_format}"

if args.w:
    subprocess.Popen([args.vlc_path, out_file, "--quiet"])

run_ffmpeg(
    [
        "-i",
        args.input_file,
        "-filter_complex_script",
        script_file,
        "-map",
        "[vout]",
        "-map",
        "[aout]",
        "-c:v",
        "libx264",
        out_file,
        "-y",
    ]
)

print(f"Done. Saved to {out_file}. Time: {datetime.now() - start_time}")


import cv2
import numpy as np
def sketch(image):
    img_gray = cv2.cvtColor(image,cv2.COLOR_BGR2GRAY)
    img_gray_blur = cv2.GaussianBlur(img_gray,(5,5),0)
    canny_edges = cv2.Canny(img_gray_blur,10,70)
    ret,mask = cv2.threshold(canny_edges,70,255,cv2.THRESH_BINARY_INV)
    return mask
cap = cv2.VideoCapture(0)
while True:
    ret,frame = cap.read()
    cv2.imshow('Live Sketch::',sketch(frame))
    if cv2.waitKey(1)==13:
        break

cap.release()
cv2.destroyAllWindows()


import os
import re
import cv2
import h5py
from bfly.logic import settings
import numpy as np
from bfly.logic import core

class DataSource(object):

    def __init__(self, core, datapath):
        '''
        '''
        self._core = core
        self._datapath = datapath
        self._size_x = -1
        self._size_y = -1
        self._size_z = -1
        # Maximum supported zoom in file system, affects calculations
        self.max_zoom = -1
        self.blocksize = (0, 0)
        self._color_map = None

    def index(self):
        '''
        Index all files without loading to
        create bounding box for this data.
        '''
        self.dtype = self.get_type()

    def get_type(self):
        return self.load(0,0,0,0).dtype

    def load_cutout(self, x0, x1, y0, y1, z, w):
        '''
        Load a cutout from a plane
        '''
        scale = float(2 ** w)
        blockshape = np.array(self.blocksize)
        x0y0 = np.array([x0,y0]) // scale
        x1y1 = np.array([x1,y1]) // scale
        top_left = np.floor(x0y0 / blockshape).astype(int)
        lo_right = np.ceil(x1y1 / blockshape).astype(int)
        origin = top_left * blockshape
        [left,top] = (x0y0 - origin).astype(int)
        [right,down] = (x1y1 - origin).astype(int)
        gridshape = lo_right-top_left
        grid = np.indices(gridshape).T

        cutshape = blockshape*grid.shape[:-1]
        cutout = np.zeros(cutshape, dtype=self.dtype)
        fixed = [z,w]
        for row in grid:
            for where in row:
                here = top_left + where
                [i0,j0] = blockshape*(where)
                [i1,j1] = blockshape*(where+1)
                one_tile = list(here)+fixed
                tile = self.load(*one_tile)
                cutout[j0:j1,i0:i1] = tile
        return cutout[top:down,left:right]

    def load(self, cur_path, w):
        '''
        Loads this file from the data path.
        '''

        cache_index = (cur_path, w)

        # Check if image in cache
        if cache_index in self._core._cache:
            # Move most recently accesed items to the top
            self._core._cache[cache_index] = self._core._cache.pop(cache_index)
            print('Loading from cache!')
            return self._core._cache[cache_index]

        # Load image from given path, check extension
        tile_ext = cur_path.rpartition('.')[2]
        if tile_ext == 'hdf5':
            with h5py.File(cur_path, 'r') as f:
                datasets = []
                f.visit(datasets.append)
                tmp_image = f[datasets[0]][()]
        else:
            print('Current path', cur_path)
            tmp_image = cv2.imread(cur_path, 0)

        # Resize if necessary, then also store to cache
        if w > 0:
            # We will use subsampling for all requests right now for speed
            if settings.ALWAYS_SUBSAMPLE or self.dtype == np.uint32:
                # Subsample to preserve accuracy for segmentations
                tmp_image = tmp_image[::2 ** w, ::2 ** w]
            else:
                factor = 0.5 ** w
                tmp_image = cv2.resize(
                    tmp_image,
                    (0,
                     0),
                    fx=factor,
                    fy=factor,
                    interpolation=settings.IMAGE_RESIZE_METHOD)

        self._core._current_cache_size += tmp_image.size

        # If we don't care to exceed the max cache size temporarily:
        # self._core._cache[cur_path] = cv2.imread(cur_path, 0)
        # self._core._current_cache_size += self._core._cache[cur_path].size

        # Remove items in cache until everything fits
        while self._core._current_cache_size > self._core._max_cache_size:
            oldest_item = list(self._core._cache.keys())[0]
            self._core._current_cache_size -= self._core._cache[
                oldest_item].size
            del self._core._cache[oldest_item]

        self._core._cache[cache_index] = tmp_image

        return tmp_image

    def seg_to_color(self, slice):
        colors = np.zeros(slice.shape+(3,),dtype=np.uint8)
        colors[:,:,0] = np.mod(107*slice[:,:],700).astype(np.uint8)
        colors[:,:,1] = np.mod(509*slice[:,:],900).astype(np.uint8)
        colors[:,:,2] = np.mod(200*slice[:,:],777).astype(np.uint8)
        return colors

    def get_boundaries(self):
        '''
        Get maximum data size
        '''
        pass

    def get_channel(self,path):
        '''
        :param path: the root/path/to/this/data
        :returns: channel of meta info about self
        '''
        dimensions = dict(list(zip(('x','y','z'),self.get_boundaries())))
        return {
            'path': path,
            'name': os.path.basename(path),
            'dimensions': dimensions,
            'data-type': str(self.dtype)
        }


import os
import tqdm
import math
import time
import argparse
import numpy as np
import tensorflow as tf
import _pickle as cPickle
from itertools import repeat
from multiprocessing import Pool
from src.official.transformer.v2 import optgen_v8, optgen_v9, optgen_v11, optgen_v12, optgen_v13, optgen_v21
from src.score_result import penalized_logp, qed, drd2, similarity
from src.official.transformer.utils.molecule_tokenizer import Moltokenizer
from src.official.transformer.v2 import optgen_v1, optgen_v2, optgen_v3, optgen_v4, optgen_v5, optgen_v6, optgen_v7

__author__ = 'Bonggun Shin'

class Timer(object):
    def __init__(self, name=None):
        self.name = name

    def __enter__(self):
        self.tstart = time.time()

    def __exit__(self, type, value, traceback):
        if self.name:
            print('[%s]' % self.name)
        print('Elapsed: %s' % (time.time() - self.tstart))

def static_var(varname, value):
    def decorate(func):
        setattr(func, varname, value)
        return func
    return decorate

@static_var("model", None)
def get_dev_model(params):
    if get_dev_model.model is None:
        # print("dev_model is created!!!")
        model_name = "optgen_%s" % params["model_version"]
        get_dev_model.model = eval(model_name).create_model(params, is_train=False)
        print("=========================== %s dev created!! ==========================" % model_name)

    else:
        print("dev_model is reused!!!")

    return get_dev_model.model

def get_layer_index(model, name=None, custom_name=None):
    for idx, l in enumerate(model.layers):
        if name=="embedding" and "embedding_freezable" in l.name:
            if custom_name==l.custom_name:
                return idx
        else:
            if name==l.name:
                return idx

    print("wrong name", name)
    exit()
    return -1

def transfer_weights_from_propnet(model, params):
    propnet_weight_path = params["propnet_weight_path"]
    print(propnet_weight_path)
    propnet_weights = cPickle.load(open(propnet_weight_path, 'rb'))

    model_layers = [l for l in model.layers]
    # model_layers[63]  # Emb 42
    # model_layers[64]  # Bidirectional 43
    # model_layers[65]  # idense
    # model_layers[66]  # Dense2 44
    layer_index = get_layer_index(model, name='embedding', custom_name='propnet_emb')
    print("setting propnet weights: %s..." % model_layers[layer_index].name)
    model_layers[layer_index].set_weights([propnet_weights["embedding_shared_weights"]])

    layer_index = get_layer_index(model, name='propnet_bidirectional')
    print("setting propnet weights: %s..." % model_layers[layer_index].name)
    lstm_weights = []
    lstm_weights.append(propnet_weights["forward_prop_lstm_0"])
    lstm_weights.append(propnet_weights["forward_prop_lstm_1"])
    lstm_weights.append(propnet_weights["forward_prop_lstm_2"])
    lstm_weights.append(propnet_weights["backward_prop_lstm_0"])
    lstm_weights.append(propnet_weights["backward_prop_lstm_1"])
    lstm_weights.append(propnet_weights["backward_prop_lstm_2"])
    model_layers[layer_index].set_weights(lstm_weights)

    layer_index = get_layer_index(model, name='propnet_idense')
    output_weights = []
    output_weights.append(propnet_weights["idense0"])
    output_weights.append(propnet_weights["idense1"])
    print("setting propnet weights: %s..." % model_layers[layer_index].name)
    model_layers[layer_index].set_weights(output_weights)

    layer_index = get_layer_index(model, name='propnet_output')
    output_weights = []
    output_weights.append(propnet_weights["output0"])
    output_weights.append(propnet_weights["output1"])
    print("setting propnet weights: %s..." % model_layers[layer_index].name)
    model_layers[layer_index].set_weights(output_weights)

    print("setting propnet weights done!!!")

def transfer_weights_from_simnet(model, params):
    simnet_weight_path = params["simnet_weight_path"]
    print(simnet_weight_path)
    simnet_weights = cPickle.load(open(simnet_weight_path, 'rb'))
    model_layers = [l for l in model.layers]
    # model_layers[63]  # Emb 42
    # model_layers[64]  # Bidirectional 43
    # model_layers[65]  # idense
    # model_layers[66]  # Dense2 44
    layer_index = get_layer_index(model, name='embedding', custom_name='simnet_emb')
    print("setting simnet weights: %s..." % model_layers[layer_index].name)
    model_layers[layer_index].set_weights([simnet_weights["embedding_shared_weights"]])

    layer_index = get_layer_index(model, name='simnet_bidirectional')
    print("setting simnet weights: %s..." % model_layers[layer_index].name)
    lstm_weights = []
    lstm_weights.append(simnet_weights["forward_sim_lstm_0"])
    lstm_weights.append(simnet_weights["forward_sim_lstm_1"])
    lstm_weights.append(simnet_weights["forward_sim_lstm_2"])
    lstm_weights.append(simnet_weights["backward_sim_lstm_0"])
    lstm_weights.append(simnet_weights["backward_sim_lstm_1"])
    lstm_weights.append(simnet_weights["backward_sim_lstm_2"])
    model_layers[layer_index].set_weights(lstm_weights)

    layer_index = get_layer_index(model, name='simnet_idense')
    output_weights = []
    output_weights.append(simnet_weights["simdense0"])
    output_weights.append(simnet_weights["simdense1"])
    print("setting simnet weights: %s..." % model_layers[layer_index].name)
    model_layers[layer_index].set_weights(output_weights)

    layer_index = get_layer_index(model, name='simnet_output')
    output_weights = []
    output_weights.append(simnet_weights["simoutput0"])
    output_weights.append(simnet_weights["simoutput1"])
    print("setting simnet weights: %s..." % model_layers[layer_index].name)
    model_layers[layer_index].set_weights(output_weights)

    print("setting simnet weights done!!!")

@static_var("model", None)
def get_trn_model(params):
    if get_trn_model.model is None:
        print("trn_model is creating!!!")
        model_name = "optgen_%s" % params["model_version"]
        get_trn_model.model = eval(model_name).create_model(params, is_train=True)
        print("=========================== %s trn created!! ==========================" % model_name)

    else:
        print("trn_model is reused!!!")

    return get_trn_model.model

def get_model(params, epoch_num):

    trn_model = get_trn_model(params)
    model_filename = "%s/%s" % (params["model_dir"], "cp-%04d.ckpt" % (epoch_num))
    print("Load weights: {}".format(model_filename))
    trn_model.load_weights(model_filename).expect_partial()

    dev_model = get_dev_model(params)
    dev_model.summary()

    # model.layers[12].set_weights(trn_model.layers[4].get_weights())
    model_name = "optgen_%s" % params["model_version"]
    trn_optgen_layer_index = get_layer_index(trn_model, name=model_name)
    dev_optgen_layer_index = get_layer_index(dev_model, name=model_name)
    dev_model.layers[dev_optgen_layer_index].set_weights(trn_model.layers[trn_optgen_layer_index].get_weights())
    print("trn_optgen_layer_index(%d, %s), dev_optgen_layer_index(%d, %s)" %
          (trn_optgen_layer_index, trn_model.layers[trn_optgen_layer_index].name,
           dev_optgen_layer_index, dev_model.layers[dev_optgen_layer_index].name))
    print("set weights for (%s) layer" % dev_model.layers[dev_optgen_layer_index].name)

    if params["use_propnet"] == 1:
        print("========================== Transferring PropNET weights....==========================")
        transfer_weights_from_propnet(dev_model, params)
    if params["use_simnet"] == 1:
        print("========================== Transferring SimNET weights....==========================")
        transfer_weights_from_simnet(dev_model, params)

    return dev_model

def get_score(src, trg, task='logp04'):
    if task=='logp04':
        logp_improvement = penalized_logp(trg) - penalized_logp(src)
        # return logp_improvement * (0.6+min(sim, 0.4))
        return logp_improvement

    elif task=='qed':
        return qed(trg)

    elif task=='drd2':
        try:
            val = drd2(trg)
        except:
            print("***************\n***************\n***************\n***************\n")
            print("***************\n***************\n***************\n***************\n")
            print(trg)
            print("***************\n***************\n***************\n***************\n")
            print("***************\n***************\n***************\n***************\n")
            val = 0
        return val

    else:
        assert 'wrong task: %s' % task

def get_dev_tst_smiles(base_path):
    file_name = '%s/simultaneous.optimization.dataset_random50.cpkl' % (base_path)

    print("============================== FILENAME ===================================")
    print(file_name)
    (selected_dev, selected_tst) = cPickle.load(open(file_name, 'rb'))

    file_name = '%s/v10.property_for_all_smiles.cpkl' % base_path
    property_dic = cPickle.load(open(file_name, 'rb'))

    return selected_dev, selected_tst, property_dic

def worker_wrapper(args):
    return worker(*args)

# def worker(return_dict, outputs, smiles_list, n, wid, start_index, end_index):
def worker(outputs, smiles_list, n, wid, start_index, end_index):
    """
    Args:
      return_dic: store result here
      n: int, number of data
      wid: int, worker id
      strat_index: int, the start index of original data that this worker will work on
      end_index: int, the end index of original data that this worker will work on
    """
    print('[worker-%d] start_index(%d), end_index(%d) n(%d)' %
          (wid, start_index, end_index, n))

    if end_index > n:
        end_index = n

    score_props = []
    for idx in range(start_index, end_index, 1):
        output = outputs[idx]
        smiles_x = smiles_list[idx]
        smiles_y = moltokenizer.decode(output)
        sim = similarity(smiles_x, smiles_y)
        if sim < 0.4:
            score_props.append([0., 0.,0.])
        else:
            score_plogp = get_score(smiles_x, smiles_y, task='logp04')
            score_qed = get_score(smiles_x, smiles_y, task='qed')
            score_drd2 = get_score(smiles_x, smiles_y, task='drd2')
            score_props.append([score_plogp, score_qed, score_drd2])

    return score_props

def evaluate(test_smiles_list, property_dic, model, moltokenizer):
    smiles_list = []
    ids_list = []
    property_x_list = []
    property_desired_list = []

    logp_list = np.array(range(3)) * 1.0 + args.lp  # 1.0 2.0 3.0
    qed_list = [0.91, 0.94, 0.97, 1.0]
    drd2_list = [0.51, 0.6, 0.7, 0.8, 0.9]

    n_grid = len(logp_list)*len(qed_list)*len(drd2_list)

    for idx, smiles in enumerate(test_smiles_list):
        # smiles = test_smiles_list[i]
        item = property_dic[smiles]
        ids = moltokenizer.encode(smiles)
        logp_val = item['logP']
        qed_val = item['qed']
        drd2_val = item['drd2']

        property_x = np.expand_dims(np.array([logp_val, qed_val, drd2_val]), axis=0)

        for logp_improvement in logp_list:
            for qed_desired in qed_list:
                for drd2_desired in drd2_list:
                    property_desired = np.expand_dims(
                        np.array([logp_val + logp_improvement, qed_desired, drd2_desired]), axis=0)

                    smiles_list.append(smiles)
                    ids_list.append(ids)
                    property_x_list.append(property_x)
                    property_desired_list.append(property_desired)

    x = np.array(tf.keras.preprocessing.sequence.pad_sequences(ids_list, dtype="int64", padding="post"))
    px = np.concatenate(property_x_list, axis=0)
    py = np.concatenate(property_desired_list, axis=0)

    outputs, _ = model.predict([x, px, py], batch_size=args.eb, verbose=1)

    if args.p == 1:
        score_plogp_all = []
        score_qed_all = []
        score_drd2_all = []
        for idx, output in tqdm.tqdm(enumerate(outputs)):
            smiles_x = smiles_list[idx]
            smiles_y = moltokenizer.decode(output)
            sim = similarity(smiles_x, smiles_y)
            if sim < 0.4:
                score_plogp_all.append(0)
                score_qed_all.append(0)
                score_drd2_all.append(0)
            else:
                score_plogp = get_score(smiles_x, smiles_y, task='logp04')
                score_qed = get_score(smiles_x, smiles_y, task='qed')
                score_drd2 = get_score(smiles_x, smiles_y, task='drd2')
                score_plogp_all.append(score_plogp)
                score_qed_all.append(score_qed)
                score_drd2_all.append(score_drd2)

    else:
        with Timer("score multi calculation..."):
            n = len(outputs)
            n_proc = args.p
            batch = math.ceil(n / (n_proc))

            with Pool(processes=n_proc) as pool:
                r = pool.map_async(worker_wrapper,
                                   zip(repeat(outputs), repeat(smiles_list),
                                       repeat(n), range(1, n_proc + 1), range(0, n, batch),
                                       range(batch, batch * n_proc + 1, batch))
                                   )
                r.wait()

            score_all = []
            for partial_score in r.get():
                score_all+=partial_score

            score_all = np.array(score_all)
            score_plogp_all = score_all[:,0]
            score_qed_all = score_all[:,1]
            score_drd2_all = score_all[:,2]

    score_plogp_final = []
    score_qed_final = []
    score_drd2_final = []
    for i in range(len(test_smiles_list)):
        score_plogp = max(score_plogp_all[n_grid * (i):n_grid * (i + 1)])
        score_qed = max(score_qed_all[n_grid * (i):n_grid * (i + 1)])
        score_drd2 = max(score_drd2_all[n_grid * (i):n_grid * (i + 1)])
        score_plogp_final.append(score_plogp)
        score_qed_final.append(score_qed)
        score_drd2_final.append(score_drd2)

    logp_mean = np.mean(score_plogp_final)
    logp_std = np.std(score_plogp_final)
    logp_rate = sum(np.array(score_plogp_final) >= 1.0) / float(len(score_plogp_final))
    qed_rate = sum(np.array(score_qed_final) >= 0.9) / float(len(score_qed_final))
    drd2_rate = sum(np.array(score_drd2_final) > 0.5) / float(len(score_drd2_final))

    all_success = 0
    for idx, logp_improvement in enumerate(score_plogp_final):
        if logp_improvement>=1.0:
            if score_qed_final[idx]>=0.9:
                if score_drd2_final[idx] > 0.5:
                    print("success index(%d/%d)" % (idx, len(score_plogp_final)))
                    all_success+=1

    all_success_rate = float(all_success)/float(len(score_plogp_final))

    return logp_mean, logp_std, logp_rate, qed_rate, drd2_rate, all_success_rate

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('-base', type=str, default="../data", help='base path.')

    parser.add_argument('-g', default="0", choices=["0", "1", "2", "3", "4", "5", "6", "7"], type=str)
    parser.add_argument('-lr', default="2", type=str) # 1e-3
    parser.add_argument('-lp', default="1.0", type=float)  # 1e-3
    parser.add_argument('-he', default=8, type=int, help='num_heads')
    parser.add_argument('-hi', default=4, type=int, help='num_hidden_layers')
    parser.add_argument('-f', default=256, type=int, help='filter_size')
    parser.add_argument('-b', default=4096, type=int, help='batch_size') # 4096, 2048
    parser.add_argument('-v', default=21, type=int, help='model version')
    parser.add_argument('-t', default=0, type=int, help='attempt')
    parser.add_argument('-es', default=1, type=int, help='eval epoch start')
    parser.add_argument('-ee', default=500, type=int, help='eval epoch end')
    parser.add_argument('-et', default=10, type=int, choices=[1, 5, 10, 20], help='eval epoch step')
    parser.add_argument('-pnet', default=1, type=int, help='if propnet used')
    parser.add_argument('-snet', default=1, type=int, help='if simnet used')
    parser.add_argument('-bf', default=1, type=int, help='if use_beam_filter')
    parser.add_argument('-p', default=32, type=int, help='number of processes')
    parser.add_argument('-eb', default=1000, type=int, help='evaluate batch')
    parser.add_argument('-rs', default=0, type=int, help='random seed')
    parser.add_argument('-ns', type=int, default=50, help='n_sample.')

    args, unparsed = parser.parse_known_args()

    os.environ["CUDA_VISIBLE_DEVICES"] = args.g
    gpus = tf.config.experimental.list_physical_devices('GPU')
    tf.config.experimental.set_memory_growth(gpus[0], True)

    model_version = "v%d" % args.v

    base_path = args.base

    n_samples = 10827615
    # dev_smiles_filename = "%s/zinc_drd2_filtered_dev_smiles_dic.cpkl" % base_path

    data_dir = "%s/v10_trn" % base_path

    iterations = 100
    steps_between_evals = n_samples // args.b
    train_steps = steps_between_evals * iterations
    validation_steps = 32
    batch_size = args.b

    model_dir = "%s/model.vv%d.vnet%d.pnet%d.snet%d.head%d.hid%d.fil%d.batch%d.lr%s.t%d" % (base_path, args.v,
                                                                                args.vnet, args.pnet,
                                                                                args.snet, args.he, args.hi,
                                                                                args.f, args.b, args.lr,
                                                                                args.t)

    mt_weight_path = "%s/mt_weights.cpkl" % base_path

    propnet_weight_path = "%s/v30.propnet.weights.cpkl" % base_path
    simnet_weight_path = "%s/v30.simnet.weights.cpkl" % base_path

    optgen_config_file = "%s/../config/optgen_config.json" % base_path

    params = optgen_v2.load_config(optgen_config_file)
    params["propnet_weight_path"] = propnet_weight_path
    params["simnet_weight_path"] = simnet_weight_path

    params["steps_between_evals"] = steps_between_evals
    params["batch_size"] = batch_size
    params["train_steps"] = train_steps
    params["validation_steps"] = validation_steps
    params["model_dir"] = model_dir
    params["mt_weight_path"] = mt_weight_path
    params["num_heads"] = args.he
    params["num_hidden_layers"] = args.hi
    params["filter_size"] = args.f
    params["learning_rate"] = float(args.lr)
    params["model_version"] = model_version

    params["n_samples"] = n_samples
    params["iterations"] = iterations
    params["data_dir"] = data_dir

    params["vocab_file"] = "%s/optgen_vocab.txt" % base_path
    params["vocab_size"] = 71

    params["use_propnet"] = args.pnet
    params["use_simnet"] = args.snet
    params["use_beam_filter"] = args.bf

    print("========================================[params]========================================")
    for k,v in params.items():
        print(k, ":" ,v)
    print("========================================[params]========================================")

    moltokenizer = Moltokenizer(params["vocab_file"])
    dev_smiles, tst_smiles, property_dic = get_dev_tst_smiles(base_path)

    epoch_list = range(args.es, args.ee, args.et)
    print("epoch_list", [i for i in epoch_list])

    dev_mean_list = []
    tst_mean_list = []

    dev_std_list = []
    tst_std_list = []

    dev_logp_rate_list = []
    tst_logp_rate_list = []

    dev_qed_list = []
    tst_qed_list = []

    dev_drd2_list = []
    tst_drd2_list = []

    dev_score_list = []
    tst_score_list = []

    dev_all_success_list = []
    tst_all_success_list = []
    for epoch_num in epoch_list:
        tst_model = get_model(params, epoch_num)

        logp_mean_dev, logp_std_dev, logp_rate_dev, qed_rate_dev, drd2_rate_dev, all_success_dev = \
            evaluate(dev_smiles, property_dic, tst_model, moltokenizer)

        logp_mean_tst, logp_std_tst, logp_rate_tst, qed_rate_tst, drd2_rate_tst, all_success_tst = \
            evaluate(tst_smiles, property_dic, tst_model, moltokenizer)

        dev_avg_score = np.mean([logp_rate_dev, qed_rate_dev, drd2_rate_dev])
        tst_avg_score = np.mean([logp_rate_tst, qed_rate_tst, drd2_rate_tst])

        dev_mean_list.append(logp_mean_dev)
        tst_mean_list.append(logp_mean_tst)

        tst_std_list.append(logp_std_tst)
        dev_std_list.append(logp_std_dev)

        dev_logp_rate_list.append(logp_rate_dev)
        tst_logp_rate_list.append(logp_rate_tst)

        dev_qed_list.append(qed_rate_dev)
        tst_qed_list.append(qed_rate_tst)

        dev_drd2_list.append(drd2_rate_dev)
        tst_drd2_list.append(drd2_rate_tst)

        dev_score_list.append(dev_avg_score)
        tst_score_list.append(tst_avg_score)

        dev_all_success_list.append(all_success_dev)
        tst_all_success_list.append(all_success_tst)

        print("======================[v%d.p%d.s%d.bf%d epoch %d]======================" %
              (args.vnet, args.pnet, args.snet, args.bf, epoch_num))
        print("[dev_logp_mean_list]")
        print(dev_mean_list)
        print("[tst_logp_mean_list]")
        print(tst_mean_list)

        print("[dev_logp_std_list]")
        print(dev_std_list)
        print("[tst_logp_std_list]")
        print(tst_std_list)

        print("[dev_logp_rate_list]")
        print(dev_logp_rate_list)
        print("[tst_logp_rate_list]")
        print(tst_logp_rate_list)

        print("[dev_qed_list]")
        print(dev_qed_list)
        print("[tst_qed_list]")
        print(tst_qed_list)

        print("[dev_drd2_list]")
        print(dev_drd2_list)
        print("[tst_drd2_list]")
        print(tst_drd2_list)

        print("[dev_score_list]")
        print(dev_score_list)
        print("[tst_score_list]")
        print(tst_score_list)

        print("dev_all_success_list")
        print(dev_all_success_list)
        print("tst_all_success_list")
        print(tst_all_success_list)

    best_index=np.argmax(dev_all_success_list)
    print("best_index(%d), best_dev(%f), best_tst(%f)" % (
    best_index, dev_all_success_list[best_index], tst_all_success_list[best_index]))

    save_path = '%s.simop.result.cpkl' % (model_dir)
    cPickle.dump((dev_mean_list, tst_mean_list, dev_std_list, tst_std_list,
                  dev_logp_rate_list, tst_logp_rate_list,
                  dev_qed_list, tst_qed_list,
                  dev_drd2_list, tst_drd2_list,
                  dev_score_list, tst_score_list,
                  dev_all_success_list, tst_all_success_list),
                 open(save_path, 'wb'))


import os
import math
import random
import numpy as np
import cv2
from PIL import Image
from collections import defaultdict

import torch
import torch.utils.data as data
import torchvision.transforms as transforms
"""
domain0 => spoof
domain1 => live
"""

class GenDataset_s(data.Dataset):
    def __init__(self, img_root, list_file, attack_type):
        super(GenDataset_s, self).__init__()

        self.img_root = img_root
        self.list_file = list_file
        self.attack_type = attack_type

        self.transform = transforms.Compose([
            transforms.ToTensor(),
        ])

        self.img_spoof_list, self.make_pair_dict = self.file_reader()

    def __getitem__(self, index):
        img_line = self.img_spoof_list[index]
        img_name_spoof, label, domain_flag, spoof_label = img_line.strip().split(' ')
        if int(spoof_label) == 2:
            spoof_type = 0
        elif int(spoof_label) == 3:
            spoof_type = 1
        elif int(spoof_label) == 4:
            spoof_type = 2
        elif int(spoof_label) == 5:
            spoof_type = 3
        else:
            print('spoof type out of except')
            spoof_type = -1
            exit()

        img_name_live = self.get_pair(label, '1')

        img_domain0 = cv2.imread(os.path.join(self.img_root, 'train_img_flod/', img_name_spoof))
        img_domain0 = cv2.cvtColor(img_domain0, cv2.COLOR_BGR2RGB)
        img_domain1 = cv2.imread(os.path.join(self.img_root, 'train_img_flod/', img_name_live))
        img_domain1 = cv2.cvtColor(img_domain1, cv2.COLOR_BGR2RGB)

        bbox_domain0_path = os.path.join(self.img_root, 'train_bbox_flod/', img_name_spoof[:-4] + '.dat')
        bbox_domain1_path = os.path.join(self.img_root, 'train_bbox_flod/', img_name_live[:-4] + '.dat')

        face_scale = 1.0

        img_domain0_crop = cv2.resize(self.crop_face_from_scene(img_domain0, bbox_domain0_path, face_scale), (256, 256))
        img_domain1_crop = cv2.resize(self.crop_face_from_scene(img_domain1, bbox_domain1_path, face_scale), (256, 256))

        img_domain0 = self.transform(img_domain0_crop)
        img_domain1 = self.transform(img_domain1_crop)

        return {'0': img_domain0, '1': img_domain1, 'type': spoof_type}

    def __len__(self):
        return len(self.img_spoof_list)

    def file_reader(self):

        def dict_profile():
            return {'0': [], '1': [], 'type': []}

        with open(self.list_file) as file:
            img_list = file.readlines()
            img_list = [x.strip() for x in img_list]

        make_pair_dict = defaultdict(dict_profile)
        img_spoof_list = []

        for line in img_list:
            img_name, label, domain_flag, spoof_label = line.strip().split(' ')
            make_pair_dict[label][domain_flag].append(img_name)

            if domain_flag == '0':
                img_spoof_list.append(line)

        return img_spoof_list, make_pair_dict

    def get_pair(self, label, domain_flag):
        img_name = random.choice(self.make_pair_dict[label][domain_flag])
        return img_name

    def crop_face_from_scene(self, image, face_name_full, scale):
        f = open(face_name_full, 'r')
        lines = f.readlines()
        lines = lines[0].split(' ')
        y1, x1, w, h = [int(ele) for ele in lines[:4]]
        f.close()
        y2 = y1 + w
        x2 = x1 + h

        y_mid = (y1 + y2) / 2.0
        x_mid = (x1 + x2) / 2.0
        h_img, w_img = image.shape[0], image.shape[1]
        # w_img,h_img=image.size
        w_scale = scale * w
        h_scale = scale * h
        y1 = y_mid - w_scale / 2.0
        x1 = x_mid - h_scale / 2.0
        y2 = y_mid + w_scale / 2.0
        x2 = x_mid + h_scale / 2.0
        y1 = max(math.floor(y1), 0)
        x1 = max(math.floor(x1), 0)
        y2 = min(math.floor(y2), w_img)
        x2 = min(math.floor(x2), h_img)

        region = image[x1:x2, y1:y2]
        return region


import numpy
# from . import txyz  # read_body and write_body live now here in mfpx.py
import logging
from molsys.util.misc import argsorted
import molsys.util.images as images

"""
These functions implement the reading and writing of mfpx files

for the handling of metadata of topologies/embedding the system of keywords is largely extended.
Note: the code still reads old mfpx file of type top correctly but it writes only the new format.

"""

logger = logging.getLogger("molsys.io")

def read(mol, f):
    """
    Read mfpx file
    :Parameters:
        -f   (obj): mfpx file object or mfpx-like readable object
        -mol (obj): instance of a molclass
    """
    ### read check ###
    try:
        f.readline ### do nothing
    except AttributeError:
        raise IOError("%s is not readable" % f)
    ### read func ###
    ftype = 'xyz'
    lbuffer = f.readline().split()
    stop = False
    bbinfo = {}
    # new dir for topo metadata with defaults
    topoinfo = {}
    topoinfo["format"]    = "old"
    topoinfo["systrekey"] = "None"
    topoinfo["RCSRname"]  = "None"
    topoinfo["spgr"]      = "None"
    topoinfo["coord_seq"] = "None"
    topoinfo["transitivity"] = "None"
    topoinfo["emapping"]  = "None"
    while not stop:
        if lbuffer[0] != '#':
            mol.natoms = int(lbuffer[0])
            stop = True
        else:
            keyword = lbuffer[1]
            if keyword == 'type':
                ftype = lbuffer[2]
            elif keyword == 'cell':
                cellparams = [float(i) for i in lbuffer[2:8]]
                mol.set_cellparams(cellparams)
            elif keyword == 'cellvect':
                mol.periodic = True
                celllist = [float(i) for i in lbuffer[2:11]]
                cell = numpy.array(celllist)
                cell.shape = (3,3)
                mol.set_cell(cell)
            elif keyword == 'bbcenter':
                mol.is_bb = True
                bbinfo['center_type'] = lbuffer[2]
                if lbuffer[2] == 'special':
                    bbinfo['center_xyz'] = numpy.array([float(i) for i in lbuffer[3:6]])
            elif keyword == 'bbconn':
                mol.is_bb = True
                con_info = lbuffer[2:]
            elif keyword == 'bbatypes':
                mol.is_bb = True
                bbinfo['connector_atypes'] = lbuffer[2:]
            elif keyword == 'orient':
                orient = [int(i) for i in lbuffer[2:]]
                mol.orientation = orient
            elif keyword[:5] == "topo_":
                assert ftype == "topo"
                topokeyw = keyword[5:]
                if topokeyw == "format":
                    if lbuffer[2] == "new":
                        topoinfo["format"] = "new"
                    else:
                        logger.warning("unknown topo format")
                elif topokeyw == "systrekey":
                    topoinfo["systrekey"] = " ".join(lbuffer[2:])
                elif topokeyw == "RCSRname":
                    topoinfo["RCSRname"] = lbuffer[2]
                elif topokeyw == "spgr":
                    topoinfo["spgr"] = lbuffer[2]
                elif topokeyw == "coord_seq":
                    topoinfo["coord_seq"] = " ".join(lbuffer[2:])
                elif topokeyw == "transitivity":
                    topoinfo["transitivity"] = " ".join(lbuffer[2:])
                elif topokeyw == "emapping":
                    topoinfo["emapping"] = [int(i) for i in lbuffer[2:]]
                else:
                    topoinfo[topokeyw.split(' ',1)[0]] = " ".join(lbuffer[2:])
                    logger.warning("unknown topo keyword %s" % keyword)
            lbuffer = f.readline().split()
    if mol.is_bb == True: ftype = 'bb'
    ### read body
    if ftype == 'xyz':
        mol.elems, mol.xyz, mol.atypes, mol.conn, mol.fragtypes, mol.fragnumbers =\
            read_body(f,mol.natoms,frags=True)
    elif ftype == 'topo':
        # topo file so set the corresponding flags
        mol.is_topo =   True
        mol.use_pconn = True
        topo_new = False
        if topoinfo["format"] == "new":
            topo_new = True
        if topo_new:
            mol.set_nofrags()
            mol.elems, mol.xyz, mol.atypes, mol.conn, mol.pconn, mol.pimages, mol.fragnumbers =\
                read_body(f,mol.natoms, topo = True, topo_new=True)
            # if we read a new topo file we add the topo addon right away and fill it with the info from topoinfo
            mol.addon("topo")
            if topoinfo["coord_seq"] == "None":
                cs = None
            else:
                cs = topoinfo["coord_seq"].split("|")
            mol.topo.topoinfo = topoinfo
            mol.topo.set_topoinfo(skey=topoinfo["systrekey"],\
                                  vmapping=mol.fragnumbers,\
                                  emapping=topoinfo["emapping"],\
                                  spgr=topoinfo["spgr"],\
                                  transitivity=topoinfo["transitivity"],\
                                  RCSRname=topoinfo["RCSRname"],\
                                  coord_seq=cs)
        else:
            mol.elems, mol.xyz, mol.atypes, mol.conn, mol.pconn, mol.pimages =\
                read_body(f,mol.natoms, topo = True)
            mol.set_nofrags()
    elif ftype == 'cromo':
        mol.is_topo =   True
        mol.is_cromo =   True
        mol.use_pconn = True
        mol.elems, mol.xyz, mol.atypes, mol.conn, mol.pconn, mol.pimages, mol.oconn =\
            read_body(f,mol.natoms,frags=True, topo = True, cromo = True)
        mol.set_nofrags()
    elif ftype == 'bb': #
        connector, connector_atoms, connector_types = parse_connstring(mol,con_info)
        #bbinfo['connector'] = connector  ## that actually has to be an argument as the setup currently is coded.
        bbinfo['connector_atoms'] = connector_atoms
        bbinfo['connector_types'] = connector_types
        mol.elems, mol.xyz, mol.atypes, mol.conn, mol.fragtypes, mol.fragnumbers =\
            read_body(f,mol.natoms,frags=True, topo = False, cromo = False)
        mol.addon('bb')
        mol.bb.setup(connector,**bbinfo)
    else:
        ftype = 'xyz'
        logger.warning('Unknown mfpx file type specified. Using xyz as default')
        mol.elems, mol.xyz, mol.atypes, mol.conn, mol.fragtypes, mol.fragnumbers =\
                read_body(f,mol.natoms,frags=False)
    mol.set_ctab_from_conn(pconn_flag=mol.use_pconn)
    mol.set_etab_from_tabs()
    if ftype == 'cromo':
        mol.set_otab_from_oconn()
    ### pass bb info
    try:
        line = f.readline().split()
        if line != [] and line[0][:5] == 'angle':
            mol.angleterm = line
    except:
        pass
    return

def write(mol, f, fullcell = True, topoformat = "new"):
    """
    Routine, which writes an mfpx file
    :Parameters:
        -mol   (obj) : instance of a molsys class
        -f (obj) : file object or writable object
        -fullcell  (bool): flag to specify if complete cellvectors arre written
    """
    ### write check ###
    try:
        f.write ### do nothing
    except AttributeError:
        raise IOError("%s is not writable" % f)
    ### write func ###
    if len(mol.fragtypes) == 0:
        mol.set_nofrags()
    if mol.is_topo:
        ftype = 'topo'
        if mol.use_pconn == False:
            # this is a topo object but without valid pconn. for writing we need to generate it
            mol.add_pconn()
        # this is for the new format
        topoinfo = {}
        if topoformat == "new":
            topoinfo["format"] = "new"
        else:
            topoinfo["format"] = "old"
        if "topo" in mol.loaded_addons:
            for k in ["systrekey", "RCSRname", "spgr", "coord_seq", "transitivity", "sk_emapping"]:
                v = getattr(mol.topo, k)
                if v != "None":
                    if type(v) == type([]):
                        if type(v[0]) == type(""):
                            # a list of strings .. use | as separator
                            topoinfo[k] = "|".join(v)
                        elif type(v[0]) == type(0):
                            # a list of integers ..
                            topoinfo[k] = " ".join([str(i) for i in v])
                    else:
                        topoinfo[k] = v
        else:
            topoinfo["format"] = "old"
    elif mol.is_bb:
        ftype = "bb"
    else:
        ftype = 'xyz'
    f.write('# type %s\n' % ftype)
    if mol.bcond>0:
        if fullcell:
#            elif keyword == 'cellvect':
#                mol.periodic = True
#                celllist = [float(i) for i in lbuffer[2:11]]
#                cell = numpy.array(celllist)
#                cell.shape = (3,3)
#                mol.set_cell(cell)
            f.write('# cellvect %12.6f %12.6f %12.6f %12.6f %12.6f %12.6f %12.6f %12.6f %12.6f\n' %\
                    tuple(mol.cell.ravel()))
        else:
            f.write('# cell %12.6f %12.6f %12.6f %12.6f %12.6f %12.6f\n' %\
                    tuple(mol.cellparams))
    if mol.is_bb:
        ### BLOCK DEFINITION ###################################################
        ### bbcenter write ###
        if mol.bb.center_type != 'special':
            f.write('# bbcenter %s\n' % mol.bb.center_type)
        else:
            f.write('# bbcenter %s %12.6f %12.6f %12.6f\n' %
                    tuple([mol.bb.center_type]+ mol.bb.center_xyz.tolist()))
        ### bbconn write ###
        # sort by connectors' types
        _argsorted_type = argsorted(mol.bb.connector_types)
        connectors_type = [mol.bb.connector_types[_] for _ in _argsorted_type]
        connector_atoms = [sorted(mol.bb.connector_atoms[_]) for _ in _argsorted_type]
        connectors      = [mol.bb.connector[_] for _ in _argsorted_type]
        # sort by connector atoms after connectors' types
        _sort_priority = [_*mol.natoms for _ in connectors_type]
        connector_atoms_priority = [
            [__+_sort_priority[_k] for __ in _] for _k,_ in enumerate(connector_atoms)
        ] # preparatory for nested sorting
        _argsorted_catoms = argsorted(connector_atoms_priority)
        connector_atoms = [connector_atoms[_] for _ in _argsorted_catoms]
        connectors      = [connectors[_]      for _ in _argsorted_catoms]
        # make bbcon line
        connstrings = ''
        ctype = 0
        for i,d in enumerate(mol.bb.connector_atoms):
            if mol.bb.connector_types[i] != ctype:
                ctype +=1
                connstrings += '/ ' # ctypes sep
            for j in d:
                connstrings = connstrings + str(j+1) +',' # connector_atoms sep
            connstrings = connstrings[0:-1] + '*' + str(connectors[i]+1)+' '
        # write bbconn line
        f.write('# bbconn %s\n' % connstrings)
        if mol.bb.connector_atypes is not None:
            # also write connector atypes
            temp_atypes = []
            for at in mol.bb.connector_atypes:
                if type(at) == type([]):
                    temp_atypes.append(",".join(at))
                else:
                    temp_atypes.append(at)
            f.write('# bbatypes %s\n' % " ".join(temp_atypes))
    if ftype == "topo":
        tks = list(topoinfo.keys())
        tks.sort()
        for k in tks:
            f.write("# topo_%s %s\n" % (k, topoinfo[k]))
    if hasattr(mol, "orientation"):
        ### ORIENTATION DEFINITION #############################################
        o = len(mol.orientation) * "%3d" % tuple(mol.orientation)
        f.write('# orient '+o+"\n")
    f.write('%i\n' % mol.natoms)
    if ftype == 'xyz' or ftype == "bb":
        write_body(f,mol)
    else:
        # this is all a hack .. restructure write/read body etc. .... currently we keep it makeing the mess even bigger .. sick!
        topo_new = False
        if topoinfo["format"] == "new":
            topo_new = True
            # make sure that fragnumbers contains the skey_mapping of the topo addon (just for storing it)
            mol.fragnumbers = mol.topo.sk_vmapping
        write_body(f,mol,topo=True, topo_new=topo_new)
    return

def parse_connstring(mol, con_info):
    """
    Routines which parses the con_info string of a txyz or an mfpx file
    :Parameters:
        - mol      (obj) : instance of a molclass
        - con_info (str) : string holding the connectors info
    """
    connector             = []
    connector_atoms       = []
    connector_types       = []
    contype_count = 0
    for c in con_info:
        if c == "/":
            contype_count += 1
        else:
            ss = c.split('*') # ss[0] is the dummy neighbors, ss[1] is the connector atom
            if len(ss) > 2:
                raise IOError('This is not a proper BB file, convert with script before!')
            elif len(ss) < 2: # neighbor == connector
                ss *= 2
            stt = ss[0].split(',')
            connector.append(int(ss[1])-1)
            connector_types.append(contype_count)
            connector_atoms.append((numpy.array([int(i) for i in stt]) -1).tolist())
    return connector, connector_atoms, connector_types

## this is read_body and write_body from txyz. due to the specific needs and special format of mfpx files these functions have been moved over here (RS)

def read_body(f, natoms, frags = True, topo = False, cromo = False, topo_new=False):
    """
    Routine, which reads the body of a txyz or a mfpx file
    :Parameters:
        -f      (obj)  : fileobject
        -natoms (int)  : number of atoms in body
        -frags  (bool) : flag to specify if fragment info is in body or not
        -topo   (bool) : flag to specify if pconn info is in body or not
        -cromo  (bool) : flag to specify if oconn info is in body or not
    """
    elems       = []
    xyz         = []
    atypes      = []
    conn        = []
    fragtypes   = []
    fragnumbers = []
    pconn       = []
    pimages     = []
    oconn      = []
    if topo: frags=False
    for i in range(natoms):
        lbuffer = f.readline().split()
        xyz.append([float(i) for i in lbuffer[2:5]])
        elems.append(lbuffer[1].lower())
        t = lbuffer[5]
        atypes.append(t)
        if frags == True:
            fragtypes.append(lbuffer[6])
            fragnumbers.append(int(lbuffer[7]))
            offset = 2
        else:
            if topo_new:
                fragtypes.append('0')  # why "0" and not "-1" ???
                fragnumbers.append(int(lbuffer[6]))
                offset = 1
            else:
                fragtypes.append('0') # why "0" and not "-1" ???
                fragnumbers.append(0)
                offset = 0
        if not topo:
            conn.append((numpy.array([int(i) for i in lbuffer[6+offset:]])-1).tolist())
        elif cromo:
            txt = lbuffer[6+offset:]
            a = [[int(j) for j in i.split('/')] for i in txt]
            c,pc,pim,oc = [i[0]-1 for i in a], [images[i[1]] for i in a], [i[1] for i in a], [i[2] for i in a]
            conn.append(c)
            pconn.append(pc)
            pimages.append(pim)
            oconn.append(oc)
        else:
            txt = lbuffer[6+offset:]
            a = [[int(j) for j in i.split('/')] for i in txt]
            c,pc,pim = [i[0]-1 for i in a], [images[i[1]] for i in a], [i[1] for i in a]
            conn.append(c)
            pconn.append(pc)
            pimages.append(pim)
    if topo:
        if cromo:
            return elems, numpy.array(xyz), atypes, conn, pconn, pimages, oconn
        else:
            if topo_new:
                return elems, numpy.array(xyz), atypes, conn, pconn, pimages, fragnumbers
            else:
#                return elems, numpy.array(xyz), atypes, conn, fragtypes, fragnumbers, pconn
                return elems, numpy.array(xyz), atypes, conn, pconn, pimages
    else:
        return elems, numpy.array(xyz), atypes, conn, fragtypes, fragnumbers

def write_body(f, mol, frags=True, topo=False, pbc=True, plain=False, topo_new=False):
    """
    Routine, which writes the body of a txyz or a mfpx file
    :Parameters:
        -f      (obj)  : fileobject
        -mol    (obj)  : instance of molsys object
        -frags  (bool) : flag to specify if fragment info should be in body or not
        -topo   (bool) : flag to specigy if pconn info should be in body or not
        -pbc    (bool) : if False, removes connectivity out of the box (meant for visualization)
        -plain  (bool) : plain Tinker file supported by molden
        -topo_new(bool): if True the new format with skey mapping (hidden in fragnumbers) is written
    """
    if topo:
        frags = False   #from now on this is convention! yuck .. why do we specify it in the first place then?
    if topo: pconn = mol.pconn
    if frags:
        fragtypes   = mol.fragtypes
        fragnumbers = mol.fragnumbers
    else:
        fragtypes = None
        if topo_new:
            fragnumbers = mol.fragnumbers ## this hides the skey mapping
        else:
            fragnumbers = None
    elems       = mol.elems
    xyz         = mol.xyz
    cnct        = mol.conn
    natoms      = mol.natoms
    atypes      = mol.atypes
    if mol.cellparams is not None and not pbc:
        mol.set_conn_nopbc()
        cnct = mol.conn_nopbc
    ### BUG but feature so removable ###
    #    atoms_withconn = mol.atoms_withconn_nopbc[:]
    #    offset = numpy.zeros(natoms, 'int')
    #    for i in range(natoms):
    #        if i not in atoms_withconn:
    #            offset[i:] += 1
    #    cnct = [
    #        [
    #            # subtract the j-th offset to atom j in i-th conn if j in atoms_withconn
    #            j-offset[j] for j in cnct[i] if j in atoms_withconn
    #        ]
    #        # if atom i in atoms_withconn
    #        for i in range(natoms) if i in atoms_withconn
    #    ]
    #    natoms = len(atoms_withconn)
    #    xyz    = xyz[atoms_withconn]
    #    elems  = numpy.take(elems, atoms_withconn).tolist()
    #    atypes = numpy.take(atypes, atoms_withconn).tolist()
    #    if frags:
    #        fragtypes   = numpy.take(fragtypes, atoms_withconn).tolist()
    #        fragnumbers = numpy.take(fragnumbers, atoms_withconn).tolist()
    if plain:
        if frags:
            fragtypes = [None]*len(atypes)
            fragnumbers = [None]*len(atypes)
            oldatypes = list(zip(atypes, fragtypes, fragnumbers))
        else:
            oldatypes = atypes[:]
        # unique atomtypes
        u_atypes = set(Counter(oldatypes).keys())
        u_atypes -= set([a for a in u_atypes if str(a).isdigit()])
        u_atypes = sorted(list(u_atypes))
        # old2new atomtypes
        o2n_atypes = {e:i for i,e in enumerate(u_atypes)}
        n2o_atypes = {i:e for i,e in enumerate(u_atypes)}
        atypes = [a if str(a).isdigit() else o2n_atypes[a] for a in oldatypes]
        frags = False ### encoded in one column only
    xyzl = xyz.tolist()
    for i in range(natoms):
        line = ("%3d %-3s" + 3*"%12.6f" + "   %-24s") % \
            (i+1, elems[i], xyzl[i][0], xyzl[i][1], xyzl[i][2], atypes[i])
        if frags is True:
            line += ("%-16s %5d ") % (fragtypes[i], fragnumbers[i])
        elif topo_new is True:
            # write fragnumbers as skey mappings if topo_new is specified
            line += ("%5d ") % (fragnumbers[i])
        conn = (numpy.array(cnct[i])+1).tolist()
        if len(conn) != 0:
            if topo:
                pimg = []
                for pc in pconn[i]:
                    for ii,img in enumerate(images):
                        if pc is None:
                            raise TypeError("Something went VERY BAD in pconn")
                        if all(img==pc):
                            pimg.append(ii)
                            break
                for cc,pp in zip(conn,pimg):
                    if pp < 10:
                        line +="%8d/%1d " % (cc,pp)
                    else:
                        line += "%7d/%2d " % (cc,pp)
            else:
                line += (len(conn)*"%7d ") % tuple(conn)
        f.write("%s \n" % line)
    if plain:
        if frags:
            f.write("### atype: (old_atype, fragment_type, fragment_number)\n")
            n2o_fmt = pprint.pformat(n2o_atypes, indent=4)
        else:
            f.write("### atype: old_atype\n")
            n2o_fmt = pprint.pformat(n2o_atypes, indent=4, width=1) #force \n
        n2o_fmt = n2o_fmt.strip("{}")
        n2o_fmt = "{\n " + n2o_fmt + "\n}\n"
        f.write(n2o_fmt)
    return


import copy

import numpy as np

from igibson.metrics.metric_base import MetricBase
from igibson.robots.manipulation_robot import IsGraspingState

class RobotMetric(MetricBase):
    def __init__(self):
        self.initialized = False

    def step_callback(self, env, _):
        robot = env.robots[0]

        self.next_state_cache = {
            "base": {"position": robot.get_position()},
        }

        self.next_state_cache.update({arm: {"position": robot.get_eef_position(arm)} for arm in robot.arm_names})

        if not self.initialized:
            self.agent_pos = {part: [] for part in ["base"] + robot.arm_names}
            self.agent_grasping = {part: [] for part in robot.arm_names}

            self.agent_local_pos = {part: [] for part in robot.arm_names}

            self.delta_agent_distance = {part: [] for part in ["base"] + robot.arm_names}
            self.delta_agent_grasp_distance = {part: [] for part in robot.arm_names}

            self.state_cache = copy.deepcopy(self.next_state_cache)
            self.initialized = True

        self.agent_pos["base"].append(list(self.state_cache["base"]["position"]))
        distance = np.linalg.norm(
            np.array(self.next_state_cache["base"]["position"]) - self.state_cache["base"]["position"]
        )
        self.delta_agent_distance["base"].append(distance)

        for arm in robot.arm_names:
            self.agent_pos[arm].append(list(self.state_cache[arm]["position"]))
            gripper_distance = np.linalg.norm(
                self.next_state_cache[arm]["position"] - self.state_cache[arm]["position"]
            )
            self.delta_agent_distance[arm].append(gripper_distance)

            self.agent_local_pos[arm].append(robot.get_eef_position(arm).tolist())

            if robot.is_grasping(arm) == IsGraspingState.TRUE:
                self.delta_agent_grasp_distance[arm].append(gripper_distance)
                self.agent_grasping[arm].append(True)
            else:
                self.delta_agent_grasp_distance[arm].append(0)
                self.agent_grasping[arm].append(False)

        self.state_cache = copy.deepcopy(self.next_state_cache)

    def gather_results(self):
        return {
            "agent_distance": {
                "timestep": self.delta_agent_distance,
            },
            "grasp_distance": {
                "timestep": self.delta_agent_grasp_distance,
            },
            "pos": {
                "timestep": self.agent_pos,
            },
            "local_pos": {
                "timestep": self.agent_local_pos,
            },
            "grasping": {
                "timestep": self.agent_grasping,
            },
        }