import os
import sqlite3 as db
import datetime
import socket
import numpy as np
import healpy as hp
import pandas as pd
import matplotlib.path as mplPath
from rubin_sim.utils import _hpid2RaDec, xyz_angular_radius, _buildTree, _xyz_from_ra_dec
from rubin_sim.site_models import FieldsDatabase
import rubin_sim

def smallest_signed_angle(a1, a2):
    """
    via https://stackoverflow.com/questions/1878907/the-smallest-difference-between-2-angles"""
    TwoPi = 2.*np.pi
    x = a1 % TwoPi
    y = a2 % TwoPi
    a = (x - y) % TwoPi
    b = (y - x) % TwoPi
    result = b+0
    alb = np.where(a < b)[0]
    result[alb] = -1.*a[alb]
    return result

class int_rounded(object):
    """
    Class to help force comparisons be made on scaled up integers,
    preventing machine precision issues cross-platforms

    Parameters
    ----------
    inval : number-like thing
        Some number that we want to compare
    scale : float (1e5)
        How much to scale inval before rounding and converting to an int.
    """
    def __init__(self, inval, scale=1e5):
        self.initial = inval
        self.value = np.round(inval * scale).astype(int)
        self.scale = scale

    def __eq__(self, other):
        return self.value == other.value

    def __ne__(self, other):
        return self.value != other.value

    def __lt__(self, other):
        return self.value < other.value

    def __le__(self, other):
        return self.value <= other.value

    def __gt__(self, other):
        return self.value > other.value

    def __ge__(self, other):
        return self.value >= other.value

    def __repr__(self):
        return str(self.initial)

    def __add__(self, other):
        out_scale = np.min([self.scale, other.scale])
        result = int_rounded(self.initial + other.initial, scale=out_scale)
        return result

    def __sub__(self, other):
        out_scale = np.min([self.scale, other.scale])
        result = int_rounded(self.initial - other.initial, scale=out_scale)
        return result

    def __mul__(self, other):
        out_scale = np.min([self.scale, other.scale])
        result = int_rounded(self.initial * other.initial, scale=out_scale)
        return result

    def __div__(self, other):
        out_scale = np.min([self.scale, other.scale])
        result = int_rounded(self.initial / other.initial, scale=out_scale)
        return result

def set_default_nside(nside=None):
    """
    Utility function to set a default nside value across the scheduler.

    XXX-there might be a better way to do this.

    Parameters
    ----------
    nside : int (None)
        A valid healpixel nside.
    """
    if not hasattr(set_default_nside, 'nside'):
        if nside is None:
            nside = 32
        set_default_nside.nside = nside
    if nside is not None:
        set_default_nside.nside = nside
    return set_default_nside.nside

def restore_scheduler(observationId, scheduler, observatory, filename, filter_sched=None):
    """Put the scheduler and observatory in the state they were in. Handy for checking reward fucnction

    Parameters
    ----------
    observationId : int
        The ID of the last observation that should be completed
    scheduler : rubin_sim.scheduler.scheduler object
        Scheduler object.
    observatory : rubin_sim.schedler.observatory.Model_observatory
        The observaotry object
    filename : str
        The output sqlite dayabase to use
    filter_sched : rubin_sim.scheduler.scheduler object
        The filter scheduler. Note that we don't look up the official end of the previous night,
        so there is potential for the loaded filters to not match.
    """
    sc = schema_converter()
    # load up the observations
    observations = sc.opsim2obs(filename)
    good_obs = np.where(observations['ID'] <= observationId)[0]
    observations = observations[good_obs]

    # replay the observations back into the scheduler
    for obs in observations:
        scheduler.add_observation(obs)
        if filter_sched is not None:
            filter_sched.add_observation(obs)

    if filter_sched is not None:
        # Make sure we have mounted the right filters for the night
        # XXX--note, this might not be exact, but should work most of the time.
        mjd_start_night = np.min(observations['mjd'][np.where(observations['night'] == obs['night'])])
        observatory.mjd = mjd_start_night
        conditions = observatory.return_conditions()
        filters_needed = filter_sched(conditions)
    else:
        filters_needed = ['u', 'g', 'r', 'i', 'y']

    # update the observatory
    observatory.mjd = obs['mjd'] + observatory.observatory.visit_time(obs)/3600./24.
    observatory.observatory.parked = False
    observatory.observatory.current_RA_rad = obs['RA']
    observatory.observatory.current_dec_rad = obs['dec']
    observatory.observatory.current_rotSkyPos_rad = obs['rotSkyPos']
    observatory.observatory.cumulative_azimuth_rad = obs['cummTelAz']
    observatory.observatory.mounted_filters = filters_needed
    # Note that we haven't updated last_az_rad, etc, but those values should be ignored.

    return scheduler, observatory

def int_binned_stat(ids, values, statistic=np.mean):
    """
    Like scipy.binned_statistic, but for unique int ids
    """

    uids = np.unique(ids)
    order = np.argsort(ids)

    ordered_ids = ids[order]
    ordered_values = values[order]

    left = np.searchsorted(ordered_ids, uids, side='left')
    right = np.searchsorted(ordered_ids, uids, side='right')

    stat_results = []
    for le, ri in zip(left, right):
        stat_results.append(statistic(ordered_values[le:ri]))

    return uids, np.array(stat_results)

def gnomonic_project_toxy(RA1, Dec1, RAcen, Deccen):
    """Calculate x/y projection of RA1/Dec1 in system with center at RAcen, Deccen.
    Input radians. Grabbed from sims_selfcal"""
    # also used in Global Telescope Network website
    cosc = np.sin(Deccen) * np.sin(Dec1) + np.cos(Deccen) * np.cos(Dec1) * np.cos(RA1-RAcen)
    x = np.cos(Dec1) * np.sin(RA1-RAcen) / cosc
    y = (np.cos(Deccen)*np.sin(Dec1) - np.sin(Deccen)*np.cos(Dec1)*np.cos(RA1-RAcen)) / cosc
    return x, y

def gnomonic_project_tosky(x, y, RAcen, Deccen):
    """Calculate RA/Dec on sky of object with x/y and RA/Cen of field of view.
    Returns Ra/Dec in radians."""
    denom = np.cos(Deccen) - y * np.sin(Deccen)
    RA = RAcen + np.arctan2(x, denom)
    Dec = np.arctan2(np.sin(Deccen) + y * np.cos(Deccen), np.sqrt(x*x + denom*denom))
    return RA, Dec

def match_hp_resolution(in_map, nside_out, UNSEEN2nan=True):
    """Utility to convert healpix map resolution if needed and change hp.UNSEEN values to
    np.nan.

    Parameters
    ----------
    in_map : np.array
        A valie healpix map
    nside_out : int
        The desired resolution to convert in_map to
    UNSEEN2nan : bool (True)
        If True, convert any hp.UNSEEN values to np.nan
    """
    current_nside = hp.npix2nside(np.size(in_map))
    if current_nside != nside_out:
        out_map = hp.ud_grade(in_map, nside_out=nside_out)
    else:
        out_map = in_map
    if UNSEEN2nan:
        out_map[np.where(out_map == hp.UNSEEN)] = np.nan
    return out_map

def raster_sort(x0, order=['x', 'y'], xbin=1.):
    """XXXX--depriciated, use tsp instead.

    Do a sort to scan a grid up and down. Simple starting guess to traveling salesman.

    Parameters
    ----------
    x0 : array
    order : list
        Keys for the order x0 should be sorted in.
    xbin : float (1.)
        The binsize to round off the first coordinate into

    returns
    -------
    array sorted so that it rasters up and down.
    """
    coords = x0.copy()
    bins = np.arange(coords[order[0]].min()-xbin/2., coords[order[0]].max()+3.*xbin/2., xbin)
    # digitize my bins
    coords[order[0]] = np.digitize(coords[order[0]], bins)
    order1 = np.argsort(coords, order=order)
    coords = coords[order1]
    places_to_invert = np.where(np.diff(coords[order[-1]]) < 0)[0]
    if np.size(places_to_invert) > 0:
        places_to_invert += 1
        indx = np.arange(coords.size)
        index_sorted = np.zeros(indx.size, dtype=int)
        index_sorted[0:places_to_invert[0]] = indx[0:places_to_invert[0]]

        for i, inv_pt in enumerate(places_to_invert[:-1]):
            if i % 2 == 0:
                index_sorted[inv_pt:places_to_invert[i+1]] = indx[inv_pt:places_to_invert[i+1]][::-1]
            else:
                index_sorted[inv_pt:places_to_invert[i+1]] = indx[inv_pt:places_to_invert[i+1]]

        if np.size(places_to_invert) % 2 != 0:
            index_sorted[places_to_invert[-1]:] = indx[places_to_invert[-1]:][::-1]
        else:
            index_sorted[places_to_invert[-1]:] = indx[places_to_invert[-1]:]
        return order1[index_sorted]
    else:
        return order1

class schema_converter(object):
    """
    Record how to convert an observation array to the standard opsim schema
    """
    def __init__(self):
        # Conversion dictionary, keys are opsim schema, values are observation dtype names
        self.convert_dict = {'observationId': 'ID', 'night': 'night',
                             'observationStartMJD': 'mjd',
                             'observationStartLST': 'lmst', 'numExposures': 'nexp',
                             'visitTime': 'visittime', 'visitExposureTime': 'exptime',
                             'proposalId': 'survey_id', 'fieldId': 'field_id',
                             'fieldRA': 'RA', 'fieldDec': 'dec', 'altitude': 'alt', 'azimuth': 'az',
                             'filter': 'filter', 'airmass': 'airmass', 'skyBrightness': 'skybrightness',
                             'cloud': 'clouds', 'seeingFwhm500': 'FWHM_500',
                             'seeingFwhmGeom': 'FWHM_geometric', 'seeingFwhmEff': 'FWHMeff',
                             'fiveSigmaDepth': 'fivesigmadepth', 'slewTime': 'slewtime',
                             'slewDistance': 'slewdist', 'paraAngle': 'pa', 'rotTelPos': 'rotTelPos',
                             'rotSkyPos': 'rotSkyPos', 'moonRA': 'moonRA',
                             'moonDec': 'moonDec', 'moonAlt': 'moonAlt', 'moonAz': 'moonAz',
                             'moonDistance': 'moonDist', 'moonPhase': 'moonPhase',
                             'sunAlt': 'sunAlt', 'sunAz': 'sunAz', 'solarElong': 'solarElong', 'note':'note'}
        # Column(s) not bothering to remap:  'observationStartTime': None,
        self.inv_map = {v: k for k, v in self.convert_dict.items()}
        # angles to converts
        self.angles_rad2deg = ['fieldRA', 'fieldDec', 'altitude', 'azimuth', 'slewDistance',
                               'paraAngle', 'rotTelPos', 'rotSkyPos', 'moonRA', 'moonDec',
                               'moonAlt', 'moonAz', 'moonDistance', 'sunAlt', 'sunAz', 'solarElong',
                               'cummTelAz']
        # Put LMST into degrees too
        self.angles_hours2deg = ['observationStartLST']

    def obs2opsim(self, obs_array, filename=None, info=None, delete_past=False):
        """convert an array of observations into a pandas dataframe with Opsim schema
        """
        if delete_past:
            try:
                os.remove(filename)
            except OSError:
                pass

        df = pd.DataFrame(obs_array)
        df = df.rename(index=str, columns=self.inv_map)
        for colname in self.angles_rad2deg:
            df[colname] = np.degrees(df[colname])
        for colname in self.angles_hours2deg:
            df[colname] = df[colname] * 360./24.

        if filename is not None:
            con = db.connect(filename)
            df.to_sql('observations', con, index=False)
            if info is not None:
                df = pd.DataFrame(info)
                df.to_sql('info', con)

    def opsim2obs(self, filename):
        """convert an opsim schema dataframe into an observation array.
        """

        con = db.connect(filename)
        df = pd.read_sql('select * from observations;', con)
        for key in self.angles_rad2deg:
            df[key] = np.radians(df[key])
        for key in self.angles_hours2deg:
            df[key] = df[key] * 24./360.

        df = df.rename(index=str, columns=self.convert_dict)

        blank = empty_observation()
        final_result = np.empty(df.shape[0], dtype=blank.dtype)
        # XXX-ugh, there has to be a better way.
        for i, key in enumerate(df.columns):
            if key in self.inv_map.keys():
                final_result[key] = df[key].values

        return final_result

def empty_observation():
    """Return a numpy array that could be a handy observation record

    XXX:  Should this really be "empty visit"? Should we have "visits" made
    up of multple "observations" to support multi-exposure time visits?

    XXX-Could add a bool flag for "observed". Then easy to track all proposed
    observations. Could also add an mjd_min, mjd_max for when an observation should be observed.
    That way we could drop things into the queue for DD fields.

    XXX--might be nice to add a generic "sched_note" str field, to record any metadata that
    would be useful to the scheduler once it's observed. and/or observationID.

    Returns
    -------
    numpy array

    The numpy fields have the following structure
    RA : float
       The Right Acension of the observation (center of the field) (Radians)
    dec : float
       Declination of the observation (Radians)
    mjd : float
       Modified Julian Date at the start of the observation (time shutter opens)
    exptime : float
       Total exposure time of the visit (seconds)
    filter : str
        The filter used. Should be one of u, g, r, i, z, y.
    rotSkyPos : float
        The rotation angle of the camera relative to the sky E of N (Radians)
    nexp : int
        Number of exposures in the visit.
    airmass : float
        Airmass at the center of the field
    FWHMeff : float
        The effective seeing FWHM at the center of the field. (arcsec)
    skybrightness : float
        The surface brightness of the sky background at the center of the
        field. (mag/sq arcsec)
    night : int
        The night number of the observation (days)
    flush_by_mjd : float
        If we hit this MJD, we should flush the queue and refill it.
    cummTelAz : float
        The cummulative telescope rotation in azimuth
    """

    names = ['ID', 'RA', 'dec', 'mjd', 'flush_by_mjd', 'exptime', 'filter', 'rotSkyPos', 'nexp',
             'airmass', 'FWHM_500', 'FWHMeff', 'FWHM_geometric', 'skybrightness', 'night',
             'slewtime', 'visittime', 'slewdist', 'fivesigmadepth',
             'alt', 'az', 'pa', 'clouds', 'moonAlt', 'sunAlt', 'note',
             'field_id', 'survey_id', 'block_id',
             'lmst', 'rotTelPos', 'moonAz', 'sunAz', 'sunRA', 'sunDec', 'moonRA', 'moonDec',
             'moonDist', 'solarElong', 'moonPhase', 'cummTelAz']

    types = [int, float, float, float, float, float, 'U1', float, int,
             float, float, float, float, float, int,
             float, float, float, float,
             float, float, float, float, float, float, 'U40',
             int, int, int,
             float, float, float, float, float, float, float, float,
             float, float, float, float]
    result = np.zeros(1, dtype=list(zip(names, types)))
    return result

def scheduled_observation():
    """Make an array for pre-scheduling observations

    mjd_tol : float
        The tolerance on how early an observation can execute (days).

    """

    # Standard things from the usual observations
    names = ['ID', 'RA', 'dec', 'mjd', 'flush_by_mjd', 'exptime', 'filter', 'rotSkyPos', 'nexp',
             'note']
    types = [int, float, float, float, float, float, 'U1', float, float, 'U40']
    names += ['mjd_tol', 'dist_tol', 'alt_min', 'alt_max', 'HA_max', 'HA_min', 'observed']
    types += [float, float, float, float, float, float, bool]
    result = np.zeros(1, dtype=list(zip(names, types)))
    return result

def read_fields():
    """Read in the Field coordinates

    Returns
    -------
    fields : `numpy.array`
        With RA and dec in radians.
    """
    query = 'select fieldId, fieldRA, fieldDEC from Field;'
    fd = FieldsDatabase()
    fields = np.array(list(fd.get_field_set(query)))
    # order by field ID
    fields = fields[fields[:,0].argsort()]

    names = ['RA', 'dec']
    types = [float, float]
    result = np.zeros(np.size(fields[:, 1]), dtype=list(zip(names, types)))
    result['RA'] = np.radians(fields[:, 1])
    result['dec'] = np.radians(fields[:, 2])

    return result

def hp_kd_tree(nside=None, leafsize=100, scale=1e5):
    """
    Generate a KD-tree of healpixel locations

    Parameters
    ----------
    nside : int
        A valid healpix nside
    leafsize : int (100)
        Leafsize of the kdtree

    Returns
    -------
    tree : scipy kdtree
    """
    if nside is None:
        nside = set_default_nside()

    hpid = np.arange(hp.nside2npix(nside))
    ra, dec = _hpid2RaDec(nside, hpid)
    return _buildTree(ra, dec, leafsize, scale=scale)

class hp_in_lsst_fov(object):
    """
    Return the healpixels within a pointing. A very simple LSST camera model with
    no chip/raft gaps.
    """
    def __init__(self, nside=None, fov_radius=1.75, scale=1e5):
        """
        Parameters
        ----------
        fov_radius : float (1.75)
            Radius of the filed of view in degrees
        """
        if nside is None:
            nside = set_default_nside()

        self.tree = hp_kd_tree(nside=nside, scale=scale)
        self.radius = np.round(xyz_angular_radius(fov_radius)*scale).astype(int)
        self.scale = scale

    def __call__(self, ra, dec, **kwargs):
        """
        Parameters
        ----------
        ra : float
            RA in radians
        dec : float
            Dec in radians

        Returns
        -------
        indx : numpy array
            The healpixels that are within the FoV
        """

        x, y, z = _xyz_from_ra_dec(np.max(ra), np.max(dec))
        x = np.round(x * self.scale).astype(int)
        y = np.round(y * self.scale).astype(int)
        z = np.round(z * self.scale).astype(int)

        indices = self.tree.query_ball_point((x, y, z), self.radius)
        return np.array(indices)

class hp_in_comcam_fov(object):
    """
    Return the healpixels within a ComCam pointing. Simple camera model
    with no chip gaps.
    """
    def __init__(self, nside=None, side_length=0.7):
        """
        Parameters
        ----------
        side_length : float (0.7)
            The length of one side of the square field of view (degrees).
        """
        if nside is None:
            nside = set_default_nside()
        self.nside = nside
        self.tree = hp_kd_tree(nside=nside)
        self.side_length = np.radians(side_length)
        self.inner_radius = xyz_angular_radius(side_length/2.)
        self.outter_radius = xyz_angular_radius(side_length/2.*np.sqrt(2.))
        # The positions of the raft corners, unrotated
        self.corners_x = np.array([-self.side_length/2., -self.side_length/2., self.side_length/2.,
                                  self.side_length/2.])
        self.corners_y = np.array([self.side_length/2., -self.side_length/2., -self.side_length/2.,
                                  self.side_length/2.])

    def __call__(self, ra, dec, rotSkyPos=0.):
        """
        Parameters
        ----------
        ra : float
            RA in radians
        dec : float
            Dec in radians
        rotSkyPos : float
            The rotation angle of the camera in radians
        Returns
        -------
        indx : numpy array
            The healpixels that are within the FoV
        """
        x, y, z = _xyz_from_ra_dec(np.max(ra), np.max(dec))
        # Healpixels within the inner circle
        indices = self.tree.query_ball_point((x, y, z), self.inner_radius)
        # Healpixels withing the outer circle
        indices_all = np.array(self.tree.query_ball_point((x, y, z), self.outter_radius))
        indices_to_check = indices_all[np.in1d(indices_all, indices, invert=True)]

        cos_rot = np.cos(rotSkyPos)
        sin_rot = np.sin(rotSkyPos)
        x_rotated = self.corners_x*cos_rot - self.corners_y*sin_rot
        y_rotated = self.corners_x*sin_rot + self.corners_y*cos_rot

        # Draw the square that we want to check if points are in.
        bbPath = mplPath.Path(np.array([[x_rotated[0], y_rotated[0]],
                                       [x_rotated[1], y_rotated[1]],
                                       [x_rotated[2], y_rotated[2]],
                                       [x_rotated[3], y_rotated[3]],
                                       [x_rotated[0], y_rotated[0]]]))

        ra_to_check, dec_to_check = _hpid2RaDec(self.nside, indices_to_check)

        # Project the indices to check to the tangent plane, see if they fall inside the polygon
        x, y = gnomonic_project_toxy(ra_to_check, dec_to_check, ra, dec)
        for i, xcheck in enumerate(x):
            # I wonder if I can do this all at once rather than a loop?
            if bbPath.contains_point((x[i], y[i])):
                indices.append(indices_to_check[i])

        return np.array(indices)

def run_info_table(observatory, extra_info=None):
    """
    Make a little table for recording the information about a run
    """

    observatory_info = observatory.get_info()
    if extra_info is not None:
        for key in extra_info:
            observatory_info.append([key, extra_info[key]])
    observatory_info = np.array(observatory_info)

    n_feature_entries = 3

    names = ['Parameter', 'Value']
    dtypes = ['|U200', '|U200']
    result = np.zeros(observatory_info[:, 0].size + n_feature_entries,
                      dtype=list(zip(names, dtypes)))

    # Fill in info about the run
    result[0]['Parameter'] = 'Date, ymd'
    now = datetime.datetime.now()
    result[0]['Value'] = '%i, %i, %i' % (now.year, now.month, now.day)

    result[1]['Parameter'] = 'hostname'
    result[1]['Value'] = socket.gethostname()

    result[2]['Parameter'] = 'rubin_sim.__version__'
    result[2]['Value'] = rubin_sim.__version__

    result[3:]['Parameter'] = observatory_info[:, 0]
    result[3:]['Value'] = observatory_info[:, 1]

    return result

def inrange(inval, minimum=-1., maximum=1.):
    """
    Make sure values are within min/max
    """
    inval = np.array(inval)
    below = np.where(inval < minimum)
    inval[below] = minimum
    above = np.where(inval > maximum)
    inval[above] = maximum
    return inval

def warm_start(scheduler, observations, mjd_key='mjd'):
    """Replay a list of observations into the scheduler

    Parameters
    ----------
    scheduler : scheduler object
    observations : np.array
        An array of observation (e.g., from sqlite2observations)
    """

    # Check that observations are in order
    observations.sort(order=mjd_key)
    for observation in observations:
        scheduler.add_observation(observation)

    return scheduler

def season_calc(night, offset=0, modulo=None, max_season=None, season_length=365.25, floor=True):
    """
    Compute what season a night is in with possible offset and modulo
    using convention that night -365 to 0 is season -1.

    Parameters
    ----------
    night : int or array
        The night we want to convert to a season
    offset : float or array (0)
        Offset to be applied to night (days)
    modulo : int (None)
        If the season should be modulated (i.e., so we can get all even years)
        (seasons, years w/default season_length)
    max_season : int (None)
        For any season above this value (before modulo), set to -1
    season_length : float (365.25)
        How long to consider one season (nights)
    floor : bool (True)
        If true, take the floor of the season. Otherwise, returns season as a float
    """
    if np.size(night) == 1:
        night = np.ravel(np.array([night]))
    result = night + offset
    result = result/season_length
    if floor:
        result = np.floor(result)
    if max_season is not None:
        over_indx = np.where(int_rounded(result) >= int_rounded(max_season))

    if modulo is not None:
        neg = np.where(int_rounded(result) < int_rounded(0))
        result = result % modulo
        result[neg] = -1
    if max_season is not None:
        result[over_indx] = -1
    if floor:
        result = result.astype(int)
    return result

def create_season_offset(nside, sun_RA_rad):
    """
    Make an offset map so seasons roll properly
    """
    hpindx = np.arange(hp.nside2npix(nside))
    ra, dec = _hpid2RaDec(nside, hpindx)
    offset = ra - sun_RA_rad + 2.*np.pi
    offset = offset % (np.pi*2)
    offset = offset * 365.25/(np.pi*2)
    offset = -offset - 365.25
    return offset

class TargetoO(object):
    """Class to hold information about a target of opportunity object

    Parameters
    ----------
    tooid : int
        Unique ID for the ToO.
    footprints : np.array
        np.array healpix maps. 1 for areas to observe, 0 for no observe.
    mjd_start : float
        The MJD the ToO starts
    duration : float
       Duration of the ToO (days).
    """
    def __init__(self, tooid, footprint, mjd_start, duration):
        self.footprint = footprint
        self.duration = duration
        self.id = tooid
        self.mjd_start = mjd_start

class Sim_targetoO_server(object):
    """Wrapper to deliver a targetoO object at the right time
    """

    def __init__(self, targetoO_list):
        self.targetoO_list = targetoO_list
        self.mjd_starts = np.array([too.mjd_start for too in self.targetoO_list])
        durations = np.array([too.duration for too in self.targetoO_list])
        self.mjd_ends = self.mjd_starts + durations

    def __call__(self, mjd):
        in_range = np.where((mjd > self.mjd_starts) & (mjd < self.mjd_ends))[0]
        result = None
        if in_range.size > 0:
            result = [self.targetoO_list[i] for i in in_range]
        return result


# Noel C. F. Codella
# Example Triplet Loss Code for Keras / TensorFlow

# Implementing Improved Triplet Loss from:
# Zhang et al. "Tracking Persons-of-Interest via Adaptive Discriminative Features" ECCV 2016

# Got help from multiple web sources, including:
# 1) https://stackoverflow.com/questions/47727679/triplet-model-for-image-retrieval-from-the-keras-pretrained-network
# 2) https://ksaluja15.github.io/Learning-Rate-Multipliers-in-Keras/
# 3) https://keras.io/preprocessing/image/
# 4) https://github.com/keras-team/keras/issues/3386
# 5) https://github.com/keras-team/keras/issues/8130

# GLOBAL DEFINES
T_G_WIDTH = 224
T_G_HEIGHT = 224
T_G_NUMCHANNELS = 3
T_G_SEED = 1337

# Misc. Necessities
import sys
import ssl # these two lines solved issues loading pretrained model
ssl._create_default_https_context = ssl._create_unverified_context
import numpy as np
import matplotlib.pyplot as plt
import cv2
from scipy.misc import imresize
np.random.seed(T_G_SEED)

# TensorFlow Includes
import tensorflow as tf
#from tensorflow.contrib.losses import metric_learning
tf.set_random_seed(T_G_SEED)

# Keras Imports & Defines
import keras
import keras.applications
from keras import backend as K
from keras.models import Model
from keras import optimizers
import keras.layers as kl

from keras.preprocessing.image import ImageDataGenerator

# Generator object for data augmentation.
# Can change values here to affect augmentation style.
datagen = ImageDataGenerator(  rotation_range=90,
                                width_shift_range=0.05,
                                height_shift_range=0.05,
                                zoom_range=0.1,
                                horizontal_flip=True,
                                vertical_flip=True,
                                )

# Local Imports
from LR_SGD import LR_SGD

# generator function for data augmentation
def createDataGen(X1, X2, X3, Y, b):

    local_seed = T_G_SEED
    genX1 = datagen.flow(X1,Y, batch_size=b, seed=local_seed, shuffle=False)
    genX2 = datagen.flow(X2,Y, batch_size=b, seed=local_seed, shuffle=False)
    genX3 = datagen.flow(X3,Y, batch_size=b, seed=local_seed, shuffle=False)
    while True:
            X1i = genX1.next()
            X2i = genX2.next()
            X3i = genX3.next()

            yield [X1i[0], X2i[0], X3i[0]], X1i[1]

def createModel(emb_size):

    # Initialize a ResNet50_ImageNet Model
    resnet_input = kl.Input(shape=(T_G_WIDTH,T_G_HEIGHT,T_G_NUMCHANNELS))
    resnet_model = keras.applications.resnet50.ResNet50(weights='imagenet', include_top = False, input_tensor=resnet_input)

    # New Layers over ResNet50
    net = resnet_model.output
    #net = kl.Flatten(name='flatten')(net)
    net = kl.GlobalAveragePooling2D(name='gap')(net)
    #net = kl.Dropout(0.5)(net)
    net = kl.Dense(emb_size,activation='relu',name='t_emb_1')(net)
    net = kl.Lambda(lambda  x: K.l2_normalize(x,axis=1), name='t_emb_1_l2norm')(net)

    # model creation
    base_model = Model(resnet_model.input, net, name="base_model")

    # triplet framework, shared weights
    input_shape=(T_G_WIDTH,T_G_HEIGHT,T_G_NUMCHANNELS)
    input_anchor = kl.Input(shape=input_shape, name='input_anchor')
    input_positive = kl.Input(shape=input_shape, name='input_pos')
    input_negative = kl.Input(shape=input_shape, name='input_neg')

    net_anchor = base_model(input_anchor)
    net_positive = base_model(input_positive)
    net_negative = base_model(input_negative)

    # The Lamda layer produces output using given function. Here its Euclidean distance.
    positive_dist = kl.Lambda(euclidean_distance, name='pos_dist')([net_anchor, net_positive])
    negative_dist = kl.Lambda(euclidean_distance, name='neg_dist')([net_anchor, net_negative])
    tertiary_dist = kl.Lambda(euclidean_distance, name='ter_dist')([net_positive, net_negative])

    # This lambda layer simply stacks outputs so both distances are available to the objective
    stacked_dists = kl.Lambda(lambda vects: K.stack(vects, axis=1), name='stacked_dists')([positive_dist, negative_dist, tertiary_dist])

    model = Model([input_anchor, input_positive, input_negative], stacked_dists, name='triple_siamese')

    # Setting up optimizer designed for variable learning rate

    # Variable Learning Rate per Layers
    lr_mult_dict = {}
    last_layer = ''
    for layer in resnet_model.layers:
        # comment this out to refine earlier layers
        # layer.trainable = False
        # print layer.name
        lr_mult_dict[layer.name] = 1
        # last_layer = layer.name
    lr_mult_dict['t_emb_1'] = 100

    base_lr = 0.0001
    momentum = 0.9
    v_optimizer = LR_SGD(lr=base_lr, momentum=momentum, decay=0.0, nesterov=False, multipliers = lr_mult_dict)

    model.compile(optimizer=v_optimizer, loss=triplet_loss, metrics=[accuracy])

    return model

def triplet_loss(y_true, y_pred):
    margin = K.constant(1)
    return K.mean(K.maximum(K.constant(0), K.square(y_pred[:,0,0]) - 0.5*(K.square(y_pred[:,1,0])+K.square(y_pred[:,2,0])) + margin))

def accuracy(y_true, y_pred):
    return K.mean(y_pred[:,0,0] < y_pred[:,1,0])

def l2Norm(x):
    return  K.l2_normalize(x, axis=-1)

def euclidean_distance(vects):
    x, y = vects
    return K.sqrt(K.maximum(K.sum(K.square(x - y), axis=1, keepdims=True), K.epsilon()))

# loads an image and preprocesses
def t_read_image(loc):
    t_image = cv2.imread(loc)
    t_image = cv2.resize(t_image, (T_G_HEIGHT,T_G_WIDTH))
    t_image = t_image.astype("float32")
    t_image = keras.applications.resnet50.preprocess_input(t_image, data_format='channels_last')

    return t_image

# loads a set of images from a text index file
def t_read_image_list(flist, start, length):

    with open(flist) as f:
        content = f.readlines()
    content = [x.strip().split()[0] for x in content]

    datalen = length
    if (datalen < 0):
        datalen = len(content)

    if (start + datalen > len(content)):
        datalen = len(content) - start

    imgset = np.zeros((datalen, T_G_HEIGHT, T_G_WIDTH, T_G_NUMCHANNELS))

    for i in range(start, start+datalen):
        if ((i-start) < len(content)):
            imgset[i-start] = t_read_image(content[i])

    return imgset

def file_numlines(fn):
    with open(fn) as f:
        return sum(1 for _ in f)

def main(argv):

    if len(argv) < 2:
        print 'Usage: \n\t -learn <Train Anchors (TXT)> <Train Positives (TXT)> <Train Negatives (TXT)> <Val Anchors (TXT)> <Val Positives (TXT)> <Val Negatives (TXT)> <embedding size> <batch size> <num epochs> <output model prefix> \n\t -extract <Model Prefix> <Input Image List (TXT)> <Output File (TXT)> \n\t\tBuilds and scores a triplet-loss model '
        return

    if 'learn' in argv[0]:
        learn(argv[1:])
    elif 'extract' in argv[0]:
        extract(argv[1:])

    return

def extract(argv):

    if len(argv) < 3:
        print 'Usage: \n\t <Model Prefix> <Input Image List (TXT)> <Output File (TXT)> \n\t\tExtracts triplet-loss model'
        return

    modelpref = argv[0]
    imglist = argv[1]
    outfile = argv[2]

    with open(modelpref + '.json', "r") as json_file:
        model_json = json_file.read()

    loaded_model = keras.models.model_from_json(model_json)
    loaded_model.load_weights(modelpref + '.h5')

    base_model = loaded_model.get_layer('base_model')

    # create a new single input
    input_shape=(T_G_WIDTH,T_G_HEIGHT,T_G_NUMCHANNELS)
    input_single = kl.Input(shape=input_shape, name='input_single')

    # create a new model without the triple loss
    net_single = base_model(input_single)
    model = Model(input_single, net_single, name='embedding_net')

    chunksize = 1000
    total_img = file_numlines(imglist)
    total_img_ch = int(np.ceil(total_img / float(chunksize)))

    with open(outfile, 'w') as f_handle:

        for i in range(0, total_img_ch):
            imgs = t_read_image_list(imglist, i*chunksize, chunksize)

            vals = model.predict(imgs)

            np.savetxt(f_handle, vals)

    return

def learn(argv):

    if len(argv) < 10:
        print 'Usage: \n\t <Train Anchors (TXT)> <Train Positives (TXT)> <Train Negatives (TXT)> <Val Anchors (TXT)> <Val Positives (TXT)> <Val Negatives (TXT)> <embedding size> <batch size> <num epochs> <output model> \n\t\tLearns triplet-loss model'
        return

    in_t_a = argv[0]
    in_t_b = argv[1]
    in_t_c = argv[2]

    in_v_a = argv[3]
    in_v_b = argv[4]
    in_v_c = argv[5]

    emb_size = int(argv[6])
    batch = int(argv[7])
    numepochs = int(argv[8])
    outpath = argv[9]

    # chunksize is the number of images we load from disk at a time
    chunksize = batch*100
    total_t = file_numlines(in_t_a)
    total_v = file_numlines(in_v_b)
    total_t_ch = int(np.ceil(total_t / float(chunksize)))
    total_v_ch = int(np.ceil(total_v / float(chunksize)))

    print 'Dataset has ' + str(total_t) + ' training triplets, and ' + str(total_v) + ' validation triplets.'

    print 'Creating a model ...'
    model = createModel(emb_size)

    print 'Training loop ...'

    # manual loop over epochs to support very large sets of triplets
    for e in range(0, numepochs):

        for t in range(0, total_t_ch):

            print 'Epoch ' + str(e) + ': train chunk ' + str(t+1) + '/ ' + str(total_t_ch) + ' ...'

            print 'Reading image lists ...'
            anchors_t = t_read_image_list(in_t_a, t*chunksize, chunksize)
            positives_t = t_read_image_list(in_t_b, t*chunksize, chunksize)
            negatives_t = t_read_image_list(in_t_c, t*chunksize, chunksize)
            Y_train = np.random.randint(2, size=(1,2,anchors_t.shape[0])).T

            print 'Starting to fit ...'
            # This method does NOT use data augmentation
            # model.fit([anchors_t, positives_t, negatives_t], Y_train, epochs=numepochs,  batch_size=batch)

            # This method uses data augmentation
            model.fit_generator(generator=createDataGen(anchors_t,positives_t,negatives_t,Y_train,batch), steps_per_epoch=len(Y_train) / batch, epochs=1, shuffle=False, use_multiprocessing=True)

        # In case the validation images don't fit in memory, we load chunks from disk again.
        val_res = [0.0, 0.0]
        total_w = 0.0
        for v in range(0, total_v_ch):

            print 'Loading validation image lists ...'
            print 'Epoch ' + str(e) + ': val chunk ' + str(v+1) + '/ ' + str(total_v_ch) + ' ...'
            anchors_v = t_read_image_list(in_v_a, v*chunksize, chunksize)
            positives_v = t_read_image_list(in_v_b, v*chunksize, chunksize)
            negatives_v = t_read_image_list(in_v_c, v*chunksize, chunksize)
            Y_val = np.random.randint(2, size=(1,2,anchors_v.shape[0])).T

            # Weight of current validation measurement.
            # if loaded expected number of items, this will be 1.0, otherwise < 1.0, and > 0.0.
            w = float(anchors_v.shape[0]) / float(chunksize)
            total_w = total_w + w

            curval = model.evaluate([anchors_v, positives_v, negatives_v], Y_val, batch_size=batch)
            val_res[0] = val_res[0] + w*curval[0]
            val_res[1] = val_res[1] + w*curval[1]

        val_res = [x / total_w for x in val_res]

        print 'Validation Results: ' + str(val_res)

    print 'Saving model ...'

    # Save the model and weights
    model.save(outpath + '.h5')

    # Due to some remaining Keras bugs around loading custom optimizers
    # and objectives, we save the model architecture as well
    model_json = model.to_json()
    with open(outpath + '.json', "w") as json_file:
        json_file.write(model_json)

    return

# Main Driver
if __name__ == "__main__":
    main(sys.argv[1:])


from itertools import chain, product
import os, sys, re, copy, ase
import ase.data as ad
from openeye.oechem import *
import numpy as np
import networkx as nx
import networkx.algorithms.isomorphism  as iso
import aqml.cheminfo.oechem.OEChem as oe
from rdkit import Chem
import scipy.spatial.distance as ssd
import aqml.cheminfo.openbabel.obabel as cib
import multiprocessing
import aqml.cheminfo.math as cim
import cml.sd as dd
import itertools as itl
import tempfile as tpf
#tsdf = tpf.NamedTemporaryFile(dir=tdir)
import aqml.cheminfo.fortran.famoneib as fa

global dsHX
dsHX = {5:1.20, 6:1.10, 7:1.00, 8:0.98, 9:0.92, 14:1.48, 15:1.42, 16:1.34, 17:1.27}

class ParentMol(object):

    def __init__(self, string, isort=False, iat=None, wg=True, k=7, \
                 k2=7, opr='.le.', fixGeom=False, keepHalogen=False, \
                 ivdw=False, dminVDW=1.2, inmr=False, \
                 covPLmin=5, debug=False):

        self.covPLmin = covPLmin

        self.k = k
        self.k2 = k2
        self.fixGeom = fixGeom
        self.iat = iat
        self.keepHalogen = keepHalogen
        self.debug = debug
        self.vsa = {'.le.': [-1,0], '.eq.': [0, ]}[opr] # valences accepted
        self.wg = wg
        self.ivdw = ivdw
        self.dminVDW = dminVDW
        self.s2cnr = {'H':1, 'B':3, 'C':4, 'N':3, 'O':2, 'F':1, \
                'Si':4, 'P':3, 'S':2, 'Cl':1, 'Br':1, 'I':1}
        self.z2cnr = {1:1, 5:3, 6:4, 7:3, 8:2, 9:1, 14:4, 15:3, 16:2, 17:1, 35:1, 53:1}

        # subg must be a subm (i.e., hybs are all retained)
        # and rings cannot be broken!
        self.FORCE_RING_CLOSED = True

        # ready the molecule
        M = oe.StringM(string, debug=debug)

        # the block below is not necessary. It's only useful
        # for test purpose to check if some fragments are missing.
        # Plus, OEChem has some limitation on num_atoms for
        # subgraph match when the subgraph is actually the whole mol
#       if not isort:
#           # protein is too huge, the program complains that
#           # the default number of matches limit reached in
#           # substructure search if you do `M.sort_atoms()
#           M.sort_atoms()

        self.M = M
        m = M.oem
        self.oem = m
        self.na = M.na
        self.g0 = ( M.bom > 0 ).astype(np.int)
        np.fill_diagonal(self.g0, 0)
        self.bom0 = M.bom

        smi = M.can
        zs0 = np.array(M.zs)
        self.zs0 = zs0
        self.coords = M.coords

        # get CNs of all heavy atoms
        cns0 = self.g0.sum(axis=0)
        self.cns0 = cns0
        cnrs0 = np.array( [ self.z2cnr[zi] for zi in zs0 ] )

        # reference net charge of each atom
        self.charges = M.charges

        # reference total valences
        self.tvs0 = self.bom0.sum(axis=0) + np.abs( self.charges )

        # get reference aromaticity of atoms
        # to be genuinely aromatic, the atom has to be unsaturated
        # Note that the so-called `genuinely aromatic means that the
        # corresponding molecular fragment cannot be described by a
        # unique SMILES string.
        ars0_1 = ( self.tvs0 - cns0 == 1 ); #print ' -- ars0_1 = ', ars0_1
        ars0_2 = [ ai.IsAromatic() for ai in m.GetAtoms() ]
        self.ars0 = np.logical_and(ars0_1, ars0_2) # ars0_1 #
        #print ' iPause, smi = ', smi
        # get envs that corresponds to multiple SMILES
        self.envs = self.get_elusive_envs()
        # get envs like 'C=C=C', 'C=C=N', 'C=N#N', etc
        self.envsC = self.get_envsC()

        if not inmr:
            ncbs = []
            if self.wg and self.ivdw:
                ncbs = M.perceive_non_covalent_bonds(dminVDW=self.dminVDW, \
                                  covPLmin=self.covPLmin)
            self.ncbs = ncbs

    def get_nodes_bridge(self, zsi, mf_i, bom_i, dsi, mapping_i_reverse, nodes_i):
        """
        get nodes connecting two or more standalone parts in a molecule/fragment
        """
        na_i = len(zsi)
        ias_i = np.arange(na_i)
        iasH = ias_i[ zsi == 1 ]
        nH = len(iasH)
        # get all pairs of H's that are not connected to the same heavy atom
        nodes_new = []
        for jh in range(nH):
            for kh in range(jh+1,nH):
                jh_u = iasH[jh]; kh_u = iasH[kh]
                h_j = mf_i.GetAtom( OEHasAtomIdx(jh_u) )
                h_k = mf_i.GetAtom( OEHasAtomIdx(kh_u) )
                nbr_jh = ias_i[ bom_i[jh_u] == 1 ][0]
                nbr_kh = ias_i[ bom_i[kh_u] == 1 ][0]
                if nbr_jh != nbr_kh:
                    dHH = dsi[kh_u,jh_u]
                    if dHH > 0 and dHH <= 1.6: # a thresh of 1.6 \AA --> ~2 heavy atoms in the shortest path will be added
                        nbr_jh_old = mapping_i_reverse[nbr_jh]
                        nbr_kh_old = mapping_i_reverse[nbr_kh]
                        a1 = self.m.GetAtom( OEHasAtomIdx(nbr_jh_old) )
                        a2 = self.m.GetAtom( OEHasAtomIdx(nbr_kh_old) )
                        #print ' nbr_jh_old, nbr_kh_old = ', nbr_jh_old, nbr_kh_old
                        for a3 in OEShortestPath(a1,a2):
                            ia3 = a3.GetIdx()
                            if ia3 not in nodes_i:
                                nodes_new.append( ia3 )
        return nodes_new

    def extend_heavy_nodes(self, jas_hvy, sg):

        degrees = self.degrees
        sets = self.sets
        ds = self.ds

        set_i = set()
        # get neighbors of those heavy atoms
        for j,ja in enumerate(jas_hvy):
            degree0_j = degrees[ja]
            degree_j = sg[j,:].sum()
            #if ja == 1339: print 'Yeah', degree_j, degree0_j
            if degree_j < degree0_j:
                if ja in self.flexible_nodes: # saturated node
                    set_i.update( [ja,] )
                else:
                    #if ja == 36: print ' Gotha 3 !'
                    for nodes_i in self.rigid_nodes:
                        if ja in nodes_i:
                            set_i.update( nodes_i ); #print ' -- ja, nodes_i = ', ja, nodes_i
            else:
                set_i.update( [ja, ] )

        jas_u = list(set_i) # nodes_of_heavy_atoms
        sets.append( set_i )
        self.sets = sets
        self.jas_u = jas_u
        #return istop

    def build_m(self, nodes_to_add):
        """
        nodes_to_add -- atomic indices to be added to build `mf
        """
        atoms = self.atoms # parent molecule

        mf = OEGraphMol()
        mapping = {}
        atoms_sg = [];

        # step 1, add heavy atoms to `mf
        icnt = 0
        for ja in nodes_to_add:
            aj = atoms[ja]; zj = self.zs0[ja]
            aj2 = mf.NewAtom( zj )
            atoms_sg.append( aj2 ); mapping[ja] = icnt; icnt += 1
            aj2.SetHyb( OEGetHybridization(aj) )
            mf.SetCoords(aj2, self.coords[ja])
        # step 2, add H's and XH bond
        bonds = []
        #print ' -- nodes_to_add = ', nodes_to_add
        for j,ja in enumerate(nodes_to_add):
            aj = atoms[ja]
            zj = self.zs0[ja]
            aj2 = atoms_sg[j]
            for ak in aj.GetAtoms():
                ka = ak.GetIdx()
                zk = ak.GetAtomicNum()
                if zk == 1:
                    ak2 = mf.NewAtom( 1 )
                    b2 = mf.NewBond( aj2, ak2, 1)
                    mf.SetCoords(ak2, self.coords[ka]); #print ' - ka, ', self.coords[ka]
                    bonds.append( [icnt,j,1] ); icnt += 1
                else:
                    # __don't__ add atom `ak to `mf as `ak may be added to `mf later
                    # in the for loop ``for ja in nodes_to_add`` later!!
                    if ka not in nodes_to_add:
                        # add H
                        v1 = self.coords[ka] - self.coords[ja]
                        dHX = dsHX[zj];
                        coords_k = self.coords[ja] + dHX*v1/np.linalg.norm(v1)
                        ak2 = mf.NewAtom( 1 )
                        mf.SetCoords(ak2, coords_k); #print ' --- ka, ', coords_k
                        b2 = mf.NewBond( aj2, ak2, 1)
                        bonds.append( [icnt,j,1] ); icnt += 1

        nadd = len(nodes_to_add)
        #print ' __ nodes_to_add = ', nodes_to_add
        for j in range(nadd):
            for k in range(j+1,nadd):
                #print ' j,k = ', j,k
                ja = nodes_to_add[j]; ka = nodes_to_add[k]
                ja2 = mapping[ja]; ka2 = mapping[ka]
                bo = self.bom0[ja,ka]
                if bo > 0:
                    aj2 = atoms_sg[ja2]; ak2 = atoms_sg[ka2]
                    bonds.append( [j,k,bo] )
                    b2 = mf.NewBond( aj2, ak2, bo )
                    #print ' (ja,ka,bo) = (%d,%d,%d), '%(ja, ka, bo), '(ja2,ka2,bo) = (%d,%d,%d)'%(ja2,ka2,bo)
        assert mf.NumAtoms() == icnt
        bom_u = np.zeros((icnt,icnt), np.int)
        for bond_i in bonds:
            bgn,end,bo_i = bond_i
            bom_u[bgn,end] = bom_u[end,bgn] = bo_i

        return bom_u, mapping, mf

    def get_rigid_and_flexible_nodes(self):
        """
        NMR only

        (1) rigid nodes
            extended smallest set of small unbreakable fragments,
            including aromatic rings, 3- and 4-membered rings
            (accompanied with high strain, not easy to cover these
            interactions in amons) and -C(=O)N- fragments

            These nodes is output as a list of lists, with each
            containing the atom indices for a unbreakable ring
            with size ranging from 3 to 9, or -C(=O)N-
        (2) flexible nodes
            a list of saturated atom indices
        """

        def update_sets(set_i, sets):
            if np.any([ set_i <= set_j for set_j in sets ]):
                return sets
            intersected = [ set_i.intersection(set_j) for set_j in sets ]
            istats = np.array([ si != set() for si in intersected ])
            nset = len(sets); idxs = np.arange( nset )
            if np.any( istats ):
                #assert istats.astype(np.int).sum() == 1
                for iset in idxs:
                    if istats[iset]:
                        sets[iset] = set_i.union( sets[iset] )
            else:
                sets.append( set_i )
            return sets

        m = self.oem
        nodes_hvy = list( np.arange(self.na)[ self.zs0 > 1 ] )

        # first search for rings
        namin = 3
        namax = 10
        sets = []
        for i in range(namin, namax+1):
            if i in [3,4,]:
                pat_i = '*~1' + '~*'*(i-2) + '~*1'
            else:
                pat_i = '*:1' + ':*'*(i-2) + ':*1'
            ss = OESubSearch(pat_i)
            iok = OEPrepareSearch(m, ss)
            for match in ss.Match(m):
                set_i = set()
                for ma in match.GetAtoms():
                    set_i.update( [ma.target.GetIdx()] )
                if set_i not in sets: sets.append( set_i )
        # now remove those rings that are union of smaller rings
        n = len(sets)
        sets_remove = []
        ijs = itl.combinations( range(n), 2 )
        sets_u = []
        for i,j in ijs:
            set_ij = sets[i].union( sets[j] )
            if set_ij in sets and (set_ij not in sets_remove):
                sets_remove.append( set_ij )
        sets_u = cim.get_compl(sets, sets_remove)
        sets = sets_u

        # then find atoms with hyb .le. 2, e.g., -C(=O)N-, -C(=O)O-,
        # -[N+](=O)[O-], -C#N, etc
        iasc = []
        for ai in m.GetAtoms():
            hyb = OEGetHybridization(ai)
            if hyb < 3 and hyb > 0:
                iasc.append( ai.GetIdx() )
        sg = self.g0[iasc,:][:,iasc]
        na_sg = len(iasc)
        dic_sg = dict(zip(range(na_sg), iasc))
        for sgi in oe.find_cliques( sg ):
            set_i = set([ dic_sg[ii] for ii in sgi ])
            sets = update_sets(set_i, sets)

        for pat_i in ['[CX3](=O)[O,N]', '[#7,#8,#9;!a][a]', ]:
            ss = OESubSearch(pat_i)
            iok = OEPrepareSearch(m, ss)
            for match in ss.Match(m):
                set_i = set()
                for ma in match.GetAtoms():
                    set_i.update( [ma.target.GetIdx()] )
                sets = update_sets(set_i, sets)
        rigid_nodes = [ list(si) for si in sets ]
        rigid_nodes_ravel = []
        for nodes_i in rigid_nodes: rigid_nodes_ravel += nodes_i
        self.rigid_nodes = rigid_nodes

        # now find flexible nodes, i.e., saturated nodes with breakable connected bonds
        flexible_nodes = list( set(nodes_hvy)^set(rigid_nodes_ravel) )

        obsolete = """
        flexible_nodes = set()
        for pat_i in ['[CX4,SiX4,PX3,F,Cl,Br,I]', ]:
            ss = OESubSearch(pat_i)
            iok = OEPrepareSearch(m, ss)
            for match in ss.Match(m):
                for ma in match.GetAtoms():
                    flexible_nodes.update( [ma.target.GetIdx()] )

        for pat_i in [ '[NX3;!a]', '[OX2;!a]', '[SX2;!a]', ]:
            ss = OESubSearch(pat_i)
            iok = OEPrepareSearch(m, ss)
            for match in ss.Match(m):
                for ma in match.GetAtoms():
                    ia = ma.target.GetIdx()
                    if ia not in rigid_nodes_ravel:
                        flexible_nodes.update( [ia] )"""

        self.flexible_nodes = list(flexible_nodes)

    def get_elusive_envs(self):
        """
        check if the bond linking two atoms in `ias is eligbile
        for breaking by inspecting if these two atoms are in a
        elusive environment, i.e., a genuinely aromatic env
        e.g., Cc1c(C)cccc1 (two SMILES exist!); exceptions:
        o1cccc1, since it has only one possible SMILES string
        """
        filt = np.array( self.ars0 ) #.astype(np.int32)
        if filt.astype(np.int).sum() == 0: return set([])
        ias0 = np.arange( len(self.ars0) )
        ias1 = ias0[filt]
        g2 = self.bom0[filt, :][:, filt]
        iok = False
        envs = set([])
        for cliques_i in oe.find_cliques(g2):
            if len(cliques_i) > 2:
                gci = g2[cliques_i, :][:, cliques_i]

                # set `irad to False to ensure that no atom is unsaturated in valence
                ess, nrss = oe.edges_standalone_updated(gci, irad=False)
                #print ' __ ess = ', ess
                #print ' __ nrss = ', nrss

                # note that for genuinely aromatic env, `nrss should contain
                # no more than 1 unique list
                nes = len(ess)
                if nes > 0:
                    n1 = len(ess[0])
                    nrs1 = set( nrss[0] )
                    for nrs in nrss[1:]:
                        if set(nrs) != nrs1:
                            raise '#ERROR: more than 2 different sets in nodes_residual??'
                #    # get the common list in `ess, then remove it
                #    # e.g., a C=C attached to a benzene ring is an
                #    # explicit env, inly the benzene ring is an implicite
                #    # env as there are more than 1 corresponding SMIELS string
                #    comms = []
                #    for i in range(n1):
                #        comm = ess[0][i]
                #        if np.all( [ comm in ess[j] for j in range(1,nes) ] ):
                #            comms.append( comm )
                #
                #    envs.update( set( oe.comba( cim.get_compl_u(ess[0],comms) ) ) )
                    envs.update( set( oe.comba( ess[0]) ) )
        envs_u = set( [ ias1[k] for k in list(envs) ] )
        return envs_u

    def get_envsC(self):
        """
        get conjugated environments containing adjacent double bonds,
        e.g., C=C=C, C=N#N
        """
        qs = ['[*]=[*]#[*]', '[*]=[*]=[*]']
        ts = []
        for q in qs:
            ots = oe.is_subg(self.oem, q, iop = 1)
            if ots[0]:
                for tsi in ots[1]:
                    tsi_u = set(tsi)
                    if len(ts) == 0:
                        ts.append( tsi_u )
                    else:
                        iexist = False
                        for j, tsj in enumerate(ts):
                            if tsi_u.intersection(tsj):
                                iexist = True
                                tsj.update( tsi_u )
                                ts[j] = tsj
                                break
                        if not iexist: ts.append( tsi_u )
        return ts

    def get_cutout(self, ias0, cutoff=8.0):
        """
        retrieve the union of local structure within a radius
        of `cutoff of atom in `ias0
        """
        m = self.oem
        self.m = m

        ias = np.arange(self.na)
        ias_hvy = ias[ self.zs0 > 1]

        self.ias_hvy = ias_hvy
        #ds = self.ds
        ds = ssd.squareform( ssd.pdist(self.coords) )
        self.ds = ds

        atoms = [ ai for ai in m.GetAtoms() ]
        self.atoms = atoms

        # get degree of heavy atom
        degrees = []
        for i in range(self.na):
            csi = self.bom0[i,:]
            degree_i = np.sum( np.logical_and( csi > 0, self.zs0 > 1 ) )
            degrees.append( degree_i )
        self.degrees = degrees

        self.get_rigid_and_flexible_nodes()

        msf = []
        self.sets = []
        boms = []
        mappings = []

        jas_u = set()
        icnt = 0
        for ia in ias0:
            filt = ( ds[ia] <= cutoff )
            jas = list( ias[filt] )

            # get heavy atoms
            jas_hvy = []
            for j,ja in enumerate(jas):
                zja = self.zs0[ja]
                if zja == 1:
                    nbr = ias[self.g0[ja,:] == 1][0]
                    #if self.zs0[nbr] in [7,8,9,15,16,17]:
                        # these electronegative atoms will induce electrostatic effects (long-ranged)
                    jas_hvy.append( nbr )
                else:
                    jas_hvy.append( ja )
            #print ' -- jas_hvy = ', jas_hvy
            # get neighbors of those heavy atoms
            sg = self.g0[jas_hvy,:][:,jas_hvy]
            #istop = self.extend_heavy_nodes(jas_hvy, sg)
            self.extend_heavy_nodes(jas_hvy, sg)
#           if 1339 in self.jas_u:
#               if icnt == 0: print self.jas_u
#               icnt += 1
            jas_u.update( self.jas_u )

        bom_u, mapping, mf = self.build_m( list(jas_u) )

        # the extracted molecular fragments (i.e., `mf) may contain
        # several disconnected components, now add some atoms
        #  re-connecting these standalone entities
        mf_i = mf
        bom_i = bom_u
        mapping_i = mapping

        mapping_i_reverse = {}
        nodes_i = [] # heavy nodes of `mf
        for keyi in mapping_i.keys():
            val_i = mapping_i[keyi]; nodes_i.append( keyi )
            mapping_i_reverse[val_i] = keyi
        if self.debug: print ' --         nodes = ', nodes_i
        dic_i = mf_i.GetCoords()
        coords_i = []
        for j in range(mf_i.NumAtoms()): coords_i.append( dic_i[j] )
        zsi = np.array([ aj.GetAtomicNum() for aj in mf_i.GetAtoms() ])
        dsi = ssd.squareform( ssd.pdist(coords_i) )

        nodes_new = self.get_nodes_bridge(zsi, mf_i, bom_i, dsi, mapping_i_reverse, nodes_i)

        if self.debug: print ' --     new nodes = ', nodes_new
        jas_hvy = list( set(nodes_i + nodes_new) )
        sg = self.g0[jas_hvy,:][:,jas_hvy]
        #istop = self.extend_heavy_nodes(jas_hvy, sg)
        #if 1339 in jas_hvy:
        #    idx = jas_hvy.index(1339)
        #    iasU = np.arange(sg.shape[0])
        #    print iasU[ sg[idx,:] > 0 ]

        self.extend_heavy_nodes(jas_hvy, sg)
        jas_u = self.jas_u

        if self.debug: print ' -- jas_u = ', jas_u, ' [updated]'
        mf_u = self.build_m( list(set(jas_u)) )[-1]
        return mf_u

    def get_atoms_within_cutoff(self, qa=None, za=None, cutoff=3.6):
        """
        For now, for prediction of NMR only

        retrieve atoms around atom `ia-th H atom within a radius of
        `cutoff.

        This function will be used when dealing with large molecules
        like proteins where long-range interactions are significant.
        The related properties include NMR shifts.
        """

        m = self.oem
        self.m = m

        ias = np.arange(self.na)
        ias_hvy = ias[ self.zs0 > 1]
        if za is None:
            ias_za = ias
        else:
            ias_za = ias[ self.zs0 == za ]
        self.ias_za = ias_za

        self.ias_hvy = ias_hvy
        #ds = self.ds
        ds = ssd.squareform( ssd.pdist(self.coords) )
        self.ds = ds

        atoms = [ ai for ai in m.GetAtoms() ]
        self.atoms = atoms

        # get degree of heavy atom
        degrees = []
        for i in range(self.na):
            csi = self.bom0[i]
            degree_i = np.sum( np.logical_and( csi > 0, self.zs0 > 1 ) )
            degrees.append( degree_i )
        self.degrees = degrees

        if qa is None:
            qsa = ias_za
        else:
            qsa = [ias_za[qa], ]
        self.get_rigid_and_flexible_nodes()

        msf = []
        self.sets = []
        boms = []
        mappings = []
        for ia in qsa:
            filt = ( ds[ia] <= cutoff )
            jas = list( ias[filt] )

            # get heavy atoms
            jas_hvy = []
            for j,ja in enumerate(jas):
                zja = self.zs0[ja]
                if zja == 1:
                    nbr = ias[self.g0[ja] == 1][0]
                    #if self.zs0[nbr] in [7,8,9,15,16,17]:
                        # these electronegative atoms will induce electrostatic effects (long-ranged)
                    jas_hvy.append( nbr )
                else:
                    jas_hvy.append( ja )
            #print ' -- jas_hvy = ', jas_hvy
            # get neighbors of those heavy atoms
            sg = self.g0[jas_hvy,:][:,jas_hvy]
            #istop = self.extend_heavy_nodes(jas_hvy, sg)
            self.extend_heavy_nodes(jas_hvy, sg)

            jas_u = self.jas_u
            #print ' -- jas_u = ', jas_u
            bom_u, mapping, mf = self.build_m(jas_u)
            boms.append(bom_u)
            mappings.append( mapping )
            msf.append(mf)

        # the extracted molecular fragments (i.e., `mf) may contain
        # several disconnected components, now add some atoms
        #  re-connecting these standalone entities
        msf_u = []
        self.sets = [] # update !! Vital!!
        for i in range(len(msf)):
            mf_i = msf[i]
            bom_i = boms[i]
            mapping_i = mappings[i]
            mapping_i_reverse = {}
            nodes_i = [] # heavy nodes of `mf
            for keyi in mapping_i.keys():
                val_i = mapping_i[keyi]; nodes_i.append( keyi )
                mapping_i_reverse[val_i] = keyi
            if self.debug: print ' --         nodes = ', nodes_i
            dic_i = mf_i.GetCoords()
            coords_i = []
            for j in range(mf_i.NumAtoms()): coords_i.append( dic_i[j] )
            zsi = np.array([ aj.GetAtomicNum() for aj in mf_i.GetAtoms() ])
            dsi = ssd.squareform( ssd.pdist(coords_i) )

            nodes_new = self.get_nodes_bridge(zsi, mf_i, bom_i, dsi, mapping_i_reverse, nodes_i)
            if self.debug: print ' --     new nodes = ', nodes_new
            jas_hvy = nodes_i + nodes_new
            sg = self.g0[jas_hvy,:][:,jas_hvy]
            #istop = self.extend_heavy_nodes(jas_hvy, sg)
            self.extend_heavy_nodes(jas_hvy, sg)
            jas_u = self.jas_u
            if self.debug: print ' -- jas_u = ', jas_u, ' [updated]'
            mf_u = self.build_m( jas_u )[-1]
            msf_u.append( mf_u )
        msf = msf_u

        # Finally remove any fragment that are part of some larger fragment
        sets = self.sets
        nmf = len(msf)
        nas = np.array( [ len(set_i) for set_i in sets ] )
        seq = np.argsort( nas )[::-1]
        #print ' -- nas = ', nas
        #print ' --seq = ', seq
        sets1 = []
        msf1 = []
        qsa1 = []
        for i in seq:
            sets1.append( sets[i ] )
            msf1.append( msf[i ] )
            qsa1.append( qsa[i ] )

        sets_u = [sets1[0], ]
        msf_u = [msf1[0], ]
        qsa_u = [ qsa1[0], ]
        for i in range( 1, nmf ):
            #print ' -- sets_u = ', sets_u
            ioks2 = [ sets1[i] <= set_j for set_j in sets_u ]
            if not np.any(ioks2): # now remove the `set_j in `sets
                sets_u.append( sets1[i] )
                msf_u.append( msf1[i] )
                qsa_u.append( qsa1[i] )

        self.sets = sets_u
        self.qsa = qsa_u
        self.msf = msf_u

class ParentMols(object):

    def __init__(self, strings, fixGeom, iat=None, wg=True, k=7,\
                 nmaxcomb=3,icc=None, substring=None, rc=6.4, \
                 isort=False, k2=7, opr='.le.', wsmi=True, irc=True, \
                 iters=[30,90], dminVDW= 1.2, \
                 idiff=0, thresh=0.2, \
                 keepHalogen=False, debug=False, ncore=1, \
                 forcefield='mmff94', do_ob_ff=True, do_rk_ff=False, \
                 ivdw=False, covPLmin=5, prefix=''):
#                 do_pm7=False, relaxHHV=False, \
        """
        prefix -- a string added to the beginning of the name of a
                  folder, where all sdf files will be written to.
                  It should be ended with '_' if it's not empty
        irc    -- T/F: relax w/wo dihedral constraints

        substring -- SMILES of a ligand.
                  Typically in a protein-ligand complex, we need
                  to identify the ligand first and then retrieve
                  all the local atoms that bind to the ligand via
                  vdW interaction as amons for training in ML. The
                  thus obtained fragment is dubbed `centre.

                  If `substring is assigned a string,
                  we will generated only amons that are
                  a) molecular complex; b) any atom in the centre
                  must be involved.
        rc     -- cutoff radius centered on each atom of the central
                  component. It's used when `icc is not None.
        """

        def check_ncbs(a, b, c):
            iok = False
            for si in itl.product(a,b):
                if set(si) in c:
                    iok = True; break
            return iok

        param = Parameters(wg, fixGeom, k, k2, ivdw, dminVDW, \
                           forcefield, thresh, do_ob_ff, \
                           do_rk_ff, idiff, iters)

        # at most 1 True can be set
        ioks = [ do_ob_ff, do_rk_ff, ]
        assert np.array(ioks).astype(np.int).sum() <= 1

        ncpu = multiprocessing.cpu_count()
        if ncore > ncpu:
            ncore = ncpu

        # temparary folder
        tdirs = ['/scratch', '/tmp']
        for tdir in tdirs:
            if os.path.exists(tdir):
                break

        # num_molecule_total
        assert type(strings) is list, '#ERROR: `strings must be a list'
        nmt = len(strings)
        if iat != None:
            assert nmt == 1, '#ERROR: if u wanna specify the atomic idx, 1 input molecule at most is allowed'

        cans = []; nhas = []; es = []; maps = []
        ms = []; ms0 = []

        # initialize `Sets
        seta = Sets(param)
        for ir in range(nmt):
            print ' -- Mid %d'%(ir+1)
            string = strings[ir]
            obj = ParentMol(string, isort=isort, iat=iat, wg=wg, k=k, k2=k2, \
                            opr=opr, fixGeom=fixGeom, covPLmin=covPLmin, \
                            ivdw=ivdw, dminVDW=dminVDW, \
                            keepHalogen=keepHalogen, debug=debug)
            ncbs = obj.ncbs
            Mlis, iass, cans = [], [], []
            # we needs all fragments in the first place; later we'll
            # remove redundencies when merging molecules to obtain
            # valid vdw complexes
            nas = []; nasv = []; pss = []
            iass = []; iassU = []
            for Mli, ias, can in obj.generate_amons():
                iasU = ias + [-1,]*(k-len(ias)); nasv.append( len(ias) )
                Mlis.append( Mli ); iass.append( ias ); cans.append( can )
                iassU.append( iasU ); pss += list(Mli[1])
                nas.append( len(Mli[0]) )
            nmi = len(cans)
            print ' -- nmi = ', nmi

            nas = np.array(nas, np.int)
            nasv = np.array(nasv, np.int)
            pss = np.array(pss)
            iassU = np.array(iassU, np.int)
            ncbsU = np.array(ncbs, np.int)

            # now combine amons to get amons complex to account for
            # long-ranged interaction
            if wg and ivdw:
                if substring != None:
                    cliques_c = set( oe.is_subg(obj.oem, substring, iop=1)[1][0] )
                    #print ' -- cliques_c = ', cliques_c
                    cliques = oe.find_cliques(obj.g0)
                    Mlis_centre = []; iass_centre = []; cans_centre = []
                    Mlis_others = []; iass_others = []; cans_others = []
                    for i in range(nmi):
                        #print ' %d/%d done'%(i+1, nmi)
                        if set(iass[i]) <= cliques_c:
                            Mlis_centre.append( Mlis[i] )
                            iass_centre.append( iass[i] )
                            cans_centre.append( cans[i] )
                        else:
                            Mlis_others.append( Mlis[i] )
                            iass_others.append( iass[i] )
                            cans_others.append( cans[i] )
                    nmi_c = len(Mlis_centre)
                    nmi_o = nmi - nmi_c
                    print ' -- nmi_centre, nmi_others = ', nmi_c, nmi_o
                    Mlis_U = []; cans_U = []
                    for i0 in range(nmi_c):
                        ias1 = iass_centre[i0]
                        t1 = Mlis_centre[i0]; nha1 = (np.array(t1[0]) > 1).sum()
                        for j0 in range(nmi_o):
                            ias2 = iass_others[j0]
                            t2 = Mlis_others[j0]; nha2 = np.array((t2[0]) > 1).sum()
                            if nha1 + nha2 <= k2 and check_ncbs(ias1, ias2, ncbs):
                                dmin = ssd.cdist(t1[1], t2[1]).min()
                                if dmin >= dminVDW:
                                    cansij = [cans_centre[i0], cans_others[j0]]
                                    cansij.sort()
                                    cans_U.append( '.'.join(cansij) )
                                    Mlis_U.append( merge(t1, t2) )
                    Mlis = Mlis_U; cans = cans_U
                    print ' -- nmi_U = ', len(Mlis)
                else:
                    obsolete = """
                    for i0 in range(nmi-1):
                        ias1 =  iass[i0]
                        t1 = Mlis[i0]; nha1 = (np.array(t1[0]) > 1).sum()
                        for j0 in range(i0+1,nmi):
                            ias2 = iass[j0]
                            t2 = Mlis[j0]; nha2 = np.array((t2[0]) > 1).sum()
                            if nha1 + nha2 <= k2 and check_ncbs(ias1, ias2, obj.ncbs):
                                dmin = ssd.cdist(t1[1], t2[1]).min()
                                if dmin >= dminVDW:
                                    cansij = [cans[i0], cans[j0]]
                                    cansij.sort()
                                    cans.append( '.'.join(cansij) )
                                    Mlis.append( merge(t1, t2) )"""
                    obsolete = """print 'nas.shape = ', nas.shape
                    print 'nasv.shape = ', nasv.shape
                    print 'iassU.shape = ', iassU.shape
                    print 'pss.shape = ', pss.shape
                    print 'ncbsU.shape = ', ncbsU.shape"""
                    print 'dminVDW = ', dminVDW
                    gv,gc = fa.get_amon_adjacency(k2,nas,nasv,iassU.T,pss.T,ncbsU.T,dminVDW)
                    print 'amon connectivity done'
                    #print 'gv=',gv # 'np.any(gv > 0) = ', np.any(gv > 0)
                    ims = np.arange(nmi)
                    combs = []
                    for im in range(nmi):
                        nv1 = nasv[im]
                        jms = ims[ gv[im] > 0 ]
                        nj = len(jms)
                        if nj == 1:
                            # in this case, nmaxcomb = 2
                            jm = jms[0]
                            if nmaxcomb == 2:
                                # setting `nmaxcomb = 2 means to include
                                # all possible combinations consisting of
                                # two standalone molecules
                                comb = [im,jms[0]]; comb.sort()
                                if comb not in combs:
                                    combs += [comb]
                            else:
                                # if we are not imposed with `nmaxcomb = 2,
                                # we remove any complex corresponding to 2) below
                                #
                                # 1)    1 --- 2  (no other frag is connected to `1 or `2)
                                #
                                # 2)    1 --- 2
                                #              \
                                #               \
                                #                3
                                if len(gv[jm]) == 1:
                                    comb = [im,jm]; comb.sort()
                                    if comb not in combs:
                                        combs += [comb]
                        else:
                            if nmaxcomb == 2:
                                for jm in jms:
                                    comb = [im,jm]; comb.sort()
                                    if comb not in combs:
                                        combs += [comb]
                            elif nmaxcomb == 3:
                                #for jm in jms:
                                #    comb = [im,jm]; comb.sort()
                                #    if comb not in combs:
                                #        combs += [comb]

                                # this is the default choice and is more reasonable
                                # as only the most relevant local frags are included.
                                # Here we don't consider frags like [im,p],[im,q] as
                                # 1) the local envs are covered by [im,p,q]; 2) it's less
                                # relevant to [im,p,q]
                                for (p,q) in itl.combinations(jms,2):
                                    nv2 = nasv[p]; nv3 = nasv[q]
                                    if nv1+nv2+nv3 <= k2 and gc[p,q] == 0:
                                        comb = [im,p,q]; comb.sort()
                                        if comb not in combs:
                                            combs += [comb]
                    print 'atom indices of all amons done'
                    for comb in combs:
                        #print comb
                        cans_i = [ cans[ic] for ic in comb ]; cans_i.sort()
                        cans.append('.'.join(cans_i))
                        ts_i = [ Mlis[ic] for ic in comb ]
                        Mlis.append( merge(ts_i) )
                    print 'amons now ready for filtering'
            #else:
            #    #

            ncan = len(cans)
            # now remove redundancy
            if wg:
                #print ' cans = ', cans
                for i in range(ncan):
                    #print '** ', cans[i], (np.array(Mlis[i][0]) > 1).sum(),\
                    #                       len(Mlis[i][0]), Mlis[i][0]
                    seta.update(ir, cans[i], Mlis[i])
                seta._sort()
            else:
                for i in range(ncan):
                    #print ' ++ i, cans[i] = ', i,cans[i]
                    seta.update2(ir, cans[i], Mlis[i])
                seta._sort2()
            print 'amons are sorted and regrouped'

        cans = seta.cans; ncs = seta.ncs; nhas = seta.nhas

        ncan = len(cans)
        self.cans = cans
        if not wsmi: return
        nd = len(str(ncan))

        s1 = 'EQ' if opr == '.eq.' else ''
        svdw = '_vdw%d'%k2 if ivdw else ''
        scomb = '_comb2' if nmaxcomb == 2 else ''
        sthresh = '_dE%.2f'%thresh if thresh > 0 else ''
        if prefix == '':
            fdn = 'g%s%d%s%s_covL%d%s'%(s1,k,svdw,sthresh,covPLmin,scomb)
        else:
            fdn = prefix

        if not os.path.exists(fdn): os.system('mkdir -p %s'%fdn)
        self.fd = fdn

        if iat is not None:
            fdn += '_iat%d'%iat # absolute idx
        if wg and (not os.path.exists(fdn+'/raw')): os.system('mkdir -p %s/raw'%fdn)
        with open(fdn + '/' + fdn+'.smi', 'w') as fid:
            fid.write('\n'.join( [ '%s %d'%(cans[i],ncs[i]) for i in range(ncan) ] ) )
        dd.io.save('%s/maps.pkl'%fdn, {'maps': maps} )

        if wg:
            ms = seta.ms; ms0 = seta.ms0;
            for i in range(ncan):
                ms_i = ms[i]; ms0_i = ms0[i]
                nci = ncs[i]
                labi = '0'*(nd - len(str(i+1))) + str(i+1)
                print ' ++ %d %06d/%06d %60s %3d'%(nhas[i], i+1, ncan, cans[i], nci)
                for j in range(nci):
                    f_j = fdn + '/frag_%s_c%05d'%(labi, j+1) + '.sdf'
                    f0_j = fdn + '/raw/frag_%s_c%05d_raw'%(labi, j+1) + '.sdf'
                    m_j = ms_i[j]; m0_j = ms0_i[j]
                    Chem.MolToMolFile(m_j, f_j)
                    Chem.MolToMolFile(m0_j, f0_j)
            print ' -- nmi_u = ', sum(ncs)
            print ' -- ncan = ', len(np.unique(cans))
        else:
            if wsmi:
                with open(fdn + '/' + fdn+'.smi', 'w') as fid:
                    fid.write('\n'.join( [ '%s'%(cans[i]) for i in range(ncan) ] ) )

import numpy,import scipy,import networkx


#!/usr/bin/env python

'''
description:    create sea surface temperatures and sea ice concentrations
                input files for Blue Action experiments 1 and 2
license:        APACHE 2.0
author:         Ronald van Haren, NLeSC (r.vanharen@esciencecenter.nl)
'''

from netCDF4 import Dataset
from netCDF4 import date2num as nc_date2num
from netCDF4 import num2date as nc_num2date
import numpy as np
import os
import time
import sys
import subprocess
import argparse

def main(startyear, endyear, overwrite):
    '''
    Main function, call experiment 1 and experiment 2 functions

    Args:
        startyear (int):        First year to use in calculations
        endyear (int):          Stop calculations in the beginning of this year
        overwrite (bool):       Force overwriting exsisting files if True
    '''
    exp1(startyear, endyear, overwrite)
    exp2(startyear, endyear, overwrite)
    exp3_4(startyear, endyear, 3, overwrite)
    exp3_4(startyear, endyear, 4, overwrite)

def create_dimensions_netcdf(ncfile, dtobj, lat, lon):
    '''
    Define dimensions in netCDF file

    Args:
        ncfile:     netCDF file handle
        dtobj:      python datetime object
        lat:        latitude [degrees]
        lon:        longitude [degrees]

    Returns:
        ncfile:     netCDF file handle
    '''
    # description of the file
    ncfile.description = 'Blue Action'
    ncfile.history = 'Created ' + time.ctime(time.time())
    # create time dimension
    timevar = ncfile.createDimension('time', None)
    # netcdf time variable UTC
    timevar = ncfile.createVariable('time', 'float32', ('time',),
                                    zlib=True)
    timevar.units = 'days since 1850-01-01 00:00:00'
    timevar.calendar = 'gregorian'
    timevar.standard_name = 'time'
    timevar.long_name = 'Time'
    timevar.axis = 'T'
    # convert dtobj to num
    timevar[:] = nc_date2num(dtobj,
                             units=ncfile['time'].units,
                             calendar=ncfile['time'].calendar)
    # write lon/lat variables
    lonvar = ncfile.createDimension('longitude', len(lon))
    lonvar = ncfile.createVariable('longitude', 'float32', ('longitude',))
    lonvar.units = 'degrees_east'
    lonvar.axis = 'X'
    lonvar.standard_name = 'longitude'
    lonvar.long_name = 'Longitude'
    lonvar[:] = lon
    latvar = ncfile.createDimension('latitude', len(lat))
    latvar = ncfile.createVariable('latitude', 'float32', ('latitude',))
    latvar.units = 'degrees_north'
    latvar.axis = 'Y'
    latvar.standard_name = 'latitude'
    latvar.long_name = 'Latitude'
    latvar[:] = lat
    return ncfile

def create_variable_sic(ncfile, sic):
    '''
    Add sea ice concentraction[time, lat, lon] to netCDF file

    Args:
        ncfile:     netCDF file handle
        sic:        sea ice concentration in [frac]

    Returns:
        ncfile:     netCDF file handle
    '''
    sicvar = ncfile.createVariable('sic', 'float32',
                                   ('time', 'latitude', 'longitude',),
                                   zlib=True, fill_value=-1e30)
    sicvar.standar_name = 'sea_ice_area_fration'
    sicvar.long_name = 'SIC'
    sicvar.units = '1'
    sicvar.cell_methods = 'time: lat: lon: mean'
    sicvar[:] = sic
    return ncfile

def create_variable_sst(ncfile, sst):
    '''
    Add sea surface temperatures[time, lat, lon] to netCDF file

    Args:
        ncfile:     netCDF file handle
        sst:        sea surface temperatures in [K]

    Returns:
        ncfile:     netCDF file handle
    '''
    sstvar = ncfile.createVariable('sst', 'float32',
                                   ('time', 'latitude', 'longitude',),
                                   fill_value=-1e30, zlib=True)
    sstvar.standar_name = 'sea_surface_temperature'
    sstvar.long_name = 'SST'
    sstvar.units = 'K'
    sstvar.cell_methods = 'time: lat: lon: mean'
    sstvar[:] = sst
    return ncfile

def sst_sic_adjustment(sst, siconc, units_sst='K', units_sic='frac'):
    '''
    Perform SST and SIC adjustments as per instructions:
        1. Set minimum SST to -1.8 degC
        2. Set SST to -1.8 degC if SIC>0.9
        3. If SST>5 degC, set SIC to 0
        4. If SIC<0.9, we calculate SSTmax, where
           SSTmax=9.328(*0.729-SIC^3)-1.8. If SST>SSTmax,
           reduce SIC so that SST=SSTmax.

    Args:
        sst:        sea surface temperatures in [degC] of [K]
        siconc:     sea ice concentration in [frac] of [perc]
        units_sst:  units of sst input, [degC] or [K]
        units_sic:  units of siconc input [frac] of [perc]

    Returns:
        sst:        adjusted sea surface temperatures in [K]
        siconc:     adjusted sea ice concentration in [frac]
    '''
    # get sst mask
    sst_mask = np.ma.getmask(sst)
    # convert units for calculation
    if (units_sst == 'K'):
        # convert to degC for calculations
        sst = sst - 273.15
    elif (units_sst == 'C'):
        pass
    else:
        print('Unknown units for SST: ' + str(units_sst))
        sys.exit()
    if (units_sic == 'perc'):
        # convert to fraction
        siconc = siconc/100.
    elif (units_sic == 'frac'):
        pass
    else:
        print('Unknown units for SIC: ' + str(units_sic))
        sys.exit()
    # set minimum SST to -1.8 degC
    np.ma.MaskedArray.clip(sst, -1.8, None, out=sst)
    # set maximum SIC to 0.999478 (IFS segfaults when SIC=1)
    np.ma.MaskedArray.clip(siconc, None, 0.999478, out=siconc)
    # unshare mask
    sst.unshare_mask()
    siconc.unshare_mask()
    # set SST to -1.8degC if sic>0.9
    sst[siconc > 0.9] = -1.8
    # If SST>5 degC, set SIC to 0
    siconc[sst > 5] = 0
    # If SIC<90%, we calculate SSTmax (SSTmax=9.328*(0.729-SIC^3)-1.8)
    sstmax = np.zeros(np.shape(siconc))
    idx_sic = (siconc < 0.9) & (siconc > 0)
    sstmax[idx_sic] = (9.328 * (0.729-(siconc[idx_sic])**3)-1.8)
    # If SST > SSTmax,reduce the SIC, so that SST=SSTmax
    idx = (sst > sstmax) & (siconc < 0.9) & (siconc > 0)
    siconc[idx] = (0.729 - ((sst[idx] + 1.8)/9.328))**(1/3)
    # set maximum SIC to 0.999478 (IFS segfaults when SIC=1)
    np.ma.MaskedArray.clip(siconc, None, 0.999478, out=siconc)
    # convert SST to K
    sst = sst + 273.15
    # reapply original sst mask
    sst = np.ma.masked_where(sst_mask, sst)
    return sst, siconc

def exp1(startyear, endyear, overwrite=False):
    '''
    create input files for Blue Action experiment 1

    Args:
        startyear (int):      First year to use in calculations
        endyear (int):        Stop calculations in the beginning of this year
        overwrite (bool):     Force overwriting existing files
    '''
    # basestring input files
    siconc_bs = ("siconc_input4MIPs_SSTsAndSeaIce_HighResMIP_MOHC-" +
                 "HadISST-2-2-0-0-0_gn_")
    tos_bs = ("tos_input4MIPs_SSTsAndSeaIce_HighResMIP_MOHC-" +
              "HadISST-2-2-0-0-0_gn_")
    # basestring output files
    tos_bsout = 'HadISST2_prelim_0to360_alldays_sst_'
    siconc_bsout = 'HadISST2_prelim_0to360_alldays_sic_'
    # create output directory if needed
    if not os.path.exists('exp1'):
        os.makedirs('exp1')
    # loop over all years
    for yr in range(int(startyear), int(endyear)):
        # input filenames
        timestr = str(yr) + '0101-' + str(yr) + '1231'
        filename_sic_in = os.path.join('siconc', siconc_bs + timestr + '.nc')
        filename_tos_in = os.path.join('tos', tos_bs + timestr + '.nc')
        # output filenames
        filename_sic_out = os.path.join('exp1', siconc_bsout + str(yr) + '.nc')
        filename_tos_out = os.path.join('exp1', tos_bsout + str(yr) + '.nc')
        # check if existing file can be used
        if (all(os.path.isfile(fl) for fl in
                [filename_sic_out, filename_tos_out]) and not overwrite):
            print('Keeping existing files: ' + filename_sic_out + ' and ' +
                  filename_tos_out)
            continue
        # open output netCDF files
        ncfile_tos_in = Dataset(filename_tos_in, 'r')
        ncfile_sic_in = Dataset(filename_sic_in, 'r')
        # open output netCDF files
        ncfile_tos = Dataset(filename_tos_out, 'w')
        ncfile_sic = Dataset(filename_sic_out, 'w')
        # convert time to datetimeobject
        dtobj = nc_num2date(ncfile_tos_in['time'][:],
                            units=ncfile_tos_in['time'].units,
                            calendar=ncfile_tos_in['time'].calendar)
        sst = ncfile_tos_in.variables['tos'][:]
        siconc = ncfile_sic_in.variables['siconc'][:]
        lat = ncfile_tos_in.variables['latitude'][:]
        lon = ncfile_tos_in.variables['longitude'][:]
        # adjust sst and sic per instructions
        sst, siconc = sst_sic_adjustment(sst, siconc,
                                         units_sst='C', units_sic='perc')
        # write sst
        ncfile_tos = create_dimensions_netcdf(ncfile_tos, dtobj, lat, lon)
        ncfile_tos = create_variable_sst(ncfile_tos, sst)
        # write sic
        ncfile_sic = create_dimensions_netcdf(ncfile_sic, dtobj, lat, lon)
        ncfile_sic = create_variable_sic(ncfile_sic, siconc)
        # close netCDF files
        ncfile_tos.close()
        ncfile_sic.close()
        ncfile_tos_in.close()
        ncfile_sic_in.close()

def exp3_4(startyear, endyear, expno, overwrite=False):
    '''
    create input files for Blue Action experiment 3 and/or 4

    Args:
        startyear (int):      First year to use in calculations
        endyear (int):        Stop calculations in the beginning of this year
        expno (int):          Experiment number (3 or 4)
        overwrite (bool):     Force overwriting existing files
    '''
    if expno not in [3, 4]:
        print('experiment number should be 3 or 4, returning...')
        return
    # basestring input files
    siconc_bs = ("siconc_input4MIPs_SSTsAndSeaIce_HighResMIP_MOHC-" +
                 "HadISST-2-2-0-0-0_gn_")
    tos_bs = ("tos_EXP" + str(expno) + "-HadISST-2-2-0-0-0_gn_")
    # basestring output files
    tos_bsout = 'HadISST2_prelim_0to360_alldays_sst_'
    siconc_bsout = 'HadISST2_prelim_0to360_alldays_sic_'
    # create output directory if needed
    if not os.path.exists('exp' + str(expno)):
        os.makedirs('exp' + str(expno))
    # loop over all years
    for yr in range(int(startyear), int(endyear)):
        # input filenames
        timestr = str(yr) + '0101-' + str(yr) + '1231'
        filename_sic_in = os.path.join('siconc', siconc_bs + timestr + '.nc')
        filename_tos_in = os.path.join('input', 'EXP' + str(expno), tos_bs + timestr + '.nc')
        # output filenames
        filename_sic_out = os.path.join('exp' + str(expno), siconc_bsout + str(yr) + '.nc')
        filename_tos_out = os.path.join('exp' + str(expno), tos_bsout + str(yr) + '.nc')
        # check if existing file can be used
        if (all(os.path.isfile(fl) for fl in
                [filename_sic_out, filename_tos_out]) and not overwrite):
            print('Keeping existing files: ' + filename_sic_out + ' and ' +
                  filename_tos_out)
            continue
        # open output netCDF files
        ncfile_tos_in = Dataset(filename_tos_in, 'r')
        ncfile_sic_in = Dataset(filename_sic_in, 'r')
        # open output netCDF files
        ncfile_tos = Dataset(filename_tos_out, 'w')
        ncfile_sic = Dataset(filename_sic_out, 'w')
        # convert time to datetimeobject
        dtobj = nc_num2date(ncfile_tos_in['time'][:],
                            units=ncfile_tos_in['time'].units,
                            calendar=ncfile_tos_in['time'].calendar)
        sst = ncfile_tos_in.variables['tos'][:]
        siconc = ncfile_sic_in.variables['siconc'][:]
        lat = ncfile_tos_in.variables['latitude'][:]
        lon = ncfile_tos_in.variables['longitude'][:]
        # adjust sst and sic per instructions
        sst, siconc = sst_sic_adjustment(sst, siconc,
                                         units_sst='C', units_sic='perc')
        # write sst
        ncfile_tos = create_dimensions_netcdf(ncfile_tos, dtobj, lat, lon)
        ncfile_tos = create_variable_sst(ncfile_tos, sst)
        # write sic
        ncfile_sic = create_dimensions_netcdf(ncfile_sic, dtobj, lat, lon)
        ncfile_sic = create_variable_sic(ncfile_sic, siconc)
        # close netCDF files
        ncfile_tos.close()
        ncfile_sic.close()
        ncfile_tos_in.close()
        ncfile_sic_in.close()

def exp2_climate(startyear, endyear, filename_sic_climate,
                 filename_sst_climate):
    '''
    Calculate daily climate for SST/SIC for Blue Action experiment 2

    Args:
        startyear (int):        First year to use in calculations
        endyear (int):          Stop calculations in the beginning of this year
        filename_sic_climate:   Path sea ice concentration climate netCDF file
        filename_sst_climate:   Path sea surface temperatures climate netCDF
                                file
    '''
    # basestring files
    sst_bs = 'HadISST2_prelim_0to360_alldays_sst_'
    siconc_bs = 'HadISST2_prelim_0to360_alldays_sic_'
    # calculate sea ice daily climate
    fn_sic = [os.path.join('exp1', siconc_bs + str(yr) + '.nc') for yr in
              range(int(startyear), int(endyear))]
    fn_sic_str = ' '.join(map(str, fn_sic))
    tmpfile = os.path.join('exp2', 'tmpfile.nc')
    command = 'cdo cat ' + fn_sic_str + ' ' + tmpfile
    subprocess.check_call(command, shell=True)
    command = 'cdo ydaymean ' + tmpfile + ' ' + filename_sic_climate
    subprocess.check_call(command, shell=True)
    os.remove(tmpfile)
    # calculate sst daily climate
    fn_sst = [os.path.join('exp1', sst_bs + str(yr) + '.nc') for yr in
              range(int(startyear), int(endyear))]
    fn_sst_str = ' '.join(map(str, fn_sst))
    command = 'cdo cat ' + fn_sst_str + ' ' + tmpfile
    subprocess.check_call(command, shell=True)
    command = 'cdo ydaymean ' + tmpfile + ' ' + filename_sst_climate
    subprocess.check_call(command, shell=True)
    os.remove(tmpfile)

def exp2(startyear, endyear, overwrite=False):
    '''
    Create input files for Blue Action experiment 2

    Args:
        startyear (int):      First year to use in calculations
        endyear (int):        Stop calculations in the beginning of this year
        overwrite (bool):     Force overwriting existing files
    '''
    # basestring files
    sst_bs = 'HadISST2_prelim_0to360_alldays_sst_'
    siconc_bs = 'HadISST2_prelim_0to360_alldays_sic_'
    # input filenames climate
    filename_sic_climate = os.path.join('exp2',
                                        'sic_climate_' + str(startyear) +
                                        '_' + str(endyear) + '.nc')
    filename_sst_climate = os.path.join('exp2',
                                        'sst_climate_' + str(startyear) +
                                        '_' + str(endyear) + '.nc')
    # create output directory if needed
    if not os.path.exists('exp2'):
        os.makedirs('exp2')
    # check if existing file can be used
    if (all(os.path.isfile(fl) for fl in
            [filename_sic_climate, filename_sst_climate]) and not overwrite):
        print('Keeping existing files: ' + filename_sic_climate +
              ' and ' + filename_sst_climate)
    else:
        exp2_climate(startyear, endyear,
                     filename_sic_climate, filename_sst_climate)
    # open sst and sic daily climate netCDF files (r)
    ncfile_sic_climate = Dataset(filename_sic_climate, 'r')
    ncfile_sst_climate = Dataset(filename_sst_climate, 'r')
    sst_clim = ncfile_sst_climate.variables['sst'][:]
    sic_clim = ncfile_sic_climate.variables['sic'][:]
    # loop over all years of exp 1 files
    for idx, yr in enumerate(range(int(startyear), int(endyear))):
        # input filenames
        filename_sic_in = os.path.join('exp1', siconc_bs + str(yr) + '.nc')
        filename_sst_in = os.path.join('exp1', sst_bs + str(yr) + '.nc')
        # output filenames
        filename_sic_out = os.path.join('exp2', siconc_bs + str(yr) + '.nc')
        filename_sst_out = os.path.join('exp2', sst_bs + str(yr) + '.nc')
        # check if existing file can be used
        if ((os.path.isfile(filename_sic_out) and
             (os.path.isfile(filename_sst_out)) and
             not overwrite)):
            print('Keeping existing files: ' + filename_sic_out +
                  ' and ' + filename_sst_out)
            continue
        # open input files (r)
        ncfile_sic_in = Dataset(filename_sic_in, 'r')
        ncfile_sst_in = Dataset(filename_sst_in, 'r')
        # open output files (rw)
        ncfile_sic_out = Dataset(filename_sic_out, 'w')
        ncfile_sst_out = Dataset(filename_sst_out, 'w')
        # convert time to datetimeobject
        dtobj = nc_num2date(ncfile_sst_in['time'][:],
                            units=ncfile_sst_in['time'].units,
                            calendar=ncfile_sst_in['time'].calendar)
        sst = ncfile_sst_in.variables['sst'][:]
        lat = ncfile_sst_in.variables['latitude'][:]
        lon = ncfile_sst_in.variables['longitude'][:]
        # get sst mask
        sst_mask = np.ma.getmask(sst)
        # use sic_clim for sic and sst_clim where sic>0
        sst = np.where(sic_clim[0:np.shape(sst)[0], :] > 0,
                       sst_clim[0:np.shape(sst)[0], :], sst)
        # reapply original sst mask
        sst = np.ma.masked_where(sst_mask, sst)
        sic = sic_clim[0:np.shape(sst)[0], :]  # make sure ndays are equal
        # redo the sst and sic adjustments as per instruction
        sst, sic = sst_sic_adjustment(sst, sic,
                                      units_sst='K', units_sic='frac')
        # write sst
        ncfile_sst_out = create_dimensions_netcdf(ncfile_sst_out, dtobj,
                                                  lat, lon)
        ncfile_sst_out = create_variable_sst(ncfile_sst_out, sst)
        # write sic
        ncfile_sic_out = create_dimensions_netcdf(ncfile_sic_out, dtobj,
                                                  lat, lon)
        ncfile_sic_out = create_variable_sic(ncfile_sic_out, sic)
        # close netCDF files
        ncfile_sic_out.close()
        ncfile_sst_out.close()
        ncfile_sic_in.close()
        ncfile_sst_in.close()
    # close daily climate netCDF files
    ncfile_sic_climate.close()
    ncfile_sst_climate.close()

if __name__ == "__main__":
    # define command line arguments
    parser = argparse.ArgumentParser()
    parser.add_argument('-s', '--startyear', type=int, default=1979,
                        help="startyear to use for processing files [default 1979]")
    parser.add_argument('-e', '--endyear', type=int, default=2016,
                        help="stop processing in the beginning of endyear [default 2016]")
    parser.add_argument('-f', '--force', action='store_true',
                        help="force overwriting existing files")
    # get arguments
    args = parser.parse_args()
    # call main()
    main(args.startyear, args.endyear, args.force)