# Licensed under a 3-clause BSD style license - see LICENSE.rst
# -*- coding: utf-8 -*-
"""
desimodel.inputs.focalplane_sync
===================================

Tools for checking and synchronizing focalplane state to online system.
"""
import os
import datetime
import shutil
import gzip
import re

import subprocess as sp

import ast

import json
import yaml

import numpy as np

from astropy.table import Table, Column

from desiutil.log import get_logger

from ..io import datadir, findfile, load_focalplane

from .focalplane_utils import (
    create_tables,
    device_loc_to_type,
    load_petal_fiber_map,
    collision_to_segments,
    exclusions_equal,
    valid_states,
    device_compare,
    device_printdiff,
    update_exclusions,
    restricted_positioner_phi,
    hash_exclusion,
)

def load_fp_calibs(path):
    """Load data dumped from the online system."""
    fpcal = Table.read(path, format="ascii.ecsv")
    return fpcal

# Rotation matrices.  These few lines of code are copied from desimeter,
# since it is the only thing that is needed from that package and there
# is nothing even desi-specific here.

def Rx(angle):  # all in radians
    Rx = np.array(
        [
            [1.0, 0.0, 0.0],
            [0.0, np.cos(angle), -np.sin(angle)],
            [0.0, np.sin(angle), np.cos(angle)],
        ]
    )
    return Rx

def Ry(angle):  # all in radians
    Ry = np.array(
        [
            [np.cos(angle), 0.0, np.sin(angle)],
            [0.0, 1.0, 0.0],
            [-np.sin(angle), 0.0, np.cos(angle)],
        ]
    )
    return Ry

def Rz(angle):  # all in radians
    Rz = np.array(
        [
            [np.cos(angle), -np.sin(angle), 0.0],
            [np.sin(angle), np.cos(angle), 0.0],
            [0.0, 0.0, 1.0],
        ]
    )
    return Rz

def Rxyz(alpha, beta, gamma):  # yaw-pitch-roll system, all in radians
    return Rz(gamma) @ Ry(beta) @ Rx(alpha)  # @ is matrix multiplication

def convert_fp_calibs(fpcal, sim=False):
    """Convert the online system information.

    This returns a tuple containing the focalplane, the current state, and the set of
    unique exclusion polygons found in the file.

    Args:
        fpcal (Table):  The table loaded from a dump from the online system.
        sim (bool):  If True, clear all transient state issues and set hardware
            to be as "good as possible", for use in simulations.

    Returns:
        (tuple):  The (focalplane, state, exclusions, time string) loaded from
            the cal file.

    """
    # Parse the calibration time
    cal_time_str = fpcal.meta["DATE_RETRIEVED"]
    cal_time = datetime.datetime.strptime(cal_time_str, "%Y-%m-%dT%H:%M:%S%z")

    state_time_str = cal_time.isoformat(timespec="seconds")

    # Parse other metadata

    eo_phi = fpcal.meta["Eo_phi"]
    eo_radius = fpcal.meta["Eo_radius_with_margin"]
    alignments = fpcal.meta["PETAL_ALIGNMENTS"]

    # Get the default exclusion polygons for theta, phi, GFA, and petal
    # boundaries

    excl = dict()

    ktheta_str = fpcal.meta["general_keepout_T"]

    ktheta_raw = np.transpose(
        np.array(ast.literal_eval(fpcal.meta["general_keepout_T"]))
    )
    kphi_raw = np.transpose(np.array(ast.literal_eval(fpcal.meta["general_keepout_P"])))
    kpetal_raw = np.transpose(np.array(ast.literal_eval(fpcal.meta["keepout_PTL"])))
    kgfa_raw = np.transpose(np.array(ast.literal_eval(fpcal.meta["keepout_GFA"])))

    kp = dict()
    kp["theta"] = dict()
    kp["theta"]["segments"] = collision_to_segments(ktheta_raw)
    kp["theta"]["circles"] = list()

    kp["phi"] = dict()
    kp["phi"]["segments"] = collision_to_segments(kphi_raw)
    kp["phi"]["circles"] = list()

    kp["petal"] = dict()
    kp["petal"]["segments"] = collision_to_segments(kpetal_raw)
    kp["petal"]["circles"] = list()

    kp["gfa"] = dict()
    kp["gfa"]["segments"] = collision_to_segments(kgfa_raw)
    kp["gfa"]["circles"] = list()

    excl["default"] = kp

    # Also make a set of exclusions for "retracted" positioners.  In this case place
    # a circle at the the center of the theta axis.

    retrct = dict(excl["default"])
    retrct["theta"]["segments"] = list()
    retrct["theta"]["circles"] = [[[0.0, 0.0], eo_radius]]
    excl["retracted"] = retrct

    # Get the fiber map from device location to spectrographs.
    fmap = load_petal_fiber_map()

    # Build the transformation information for each petal alignment
    petal_rot = dict()
    petal_trans = dict()
    for petal_id, align in alignments.items():
        petal_rot[petal_id] = Rxyz(align["alpha"], align["beta"], align["gamma"])
        petal_trans[petal_id] = np.array([align["Tx"], align["Ty"], align["Tz"]])

    n_rows = len(fpcal)

    fp, state = create_tables(n_rows)

    # We only want to track the POS_P and POS_T values for positioners which are
    # non-functional (stuck) or have a broken fiber, since these are the positioners
    # we cannot move.
    stuck_or_broken = valid_states["BROKEN"] | valid_states["STUCK"]

    kindx = 0

    for r in range(n_rows):
        d = fpcal[r]
        # First set the focalplane properties
        fp["PETAL"][r] = d["PETAL_LOC"]
        fp["PETAL_ID"][r] = d["PETAL_ID"]
        fp["DEVICE"][r] = d["DEVICE_LOC"]
        fp["DEVICE_ID"][r] = d["POS_ID"]
        fp["LOCATION"][r] = d["PETAL_LOC"] * 1000 + d["DEVICE_LOC"]
        fp["DEVICE_TYPE"][r] = device_loc_to_type(fp["DEVICE"][r])
        fp["SLITBLOCK"][r] = -1
        fp["BLOCKFIBER"][r] = -1
        fp["CABLE"][r] = -1
        fp["CONDUIT"][r] = "NA"
        fp["FWHM"][r] = 0.0
        fp["FRD"][r] = 0.0
        fp["ABS"][r] = 0.0
        if d["PETAL_ID"] in fmap:
            # We have some information about the device to fiber mapping.
            if d["DEVICE_LOC"] in fmap[d["PETAL_ID"]]:
                # This is a POS or ETC device
                fp["SLITBLOCK"][r] = fmap[d["PETAL_ID"]][d["DEVICE_LOC"]]["SLITBLOCK"]
                fp["BLOCKFIBER"][r] = fmap[d["PETAL_ID"]][d["DEVICE_LOC"]]["BLOCKFIBER"]
                fp["CABLE"][r] = fmap[d["PETAL_ID"]][d["DEVICE_LOC"]]["CABLE"]
                fp["CONDUIT"][r] = fmap[d["PETAL_ID"]][d["DEVICE_LOC"]]["CONDUIT"]
                fp["FWHM"][r] = fmap[d["PETAL_ID"]][d["DEVICE_LOC"]]["FWHM"]
                fp["FRD"][r] = fmap[d["PETAL_ID"]][d["DEVICE_LOC"]]["FRD"]
                fp["ABS"][r] = fmap[d["PETAL_ID"]][d["DEVICE_LOC"]]["ABS"]
        fp["LENGTH_R1"][r] = d["LENGTH_R1"]
        fp["LENGTH_R2"][r] = d["LENGTH_R2"]
        fp["OFFSET_X"][r] = d["OFFSET_X_CS5"]
        fp["OFFSET_Y"][r] = d["OFFSET_Y_CS5"]
        fp["OFFSET_P"][r] = d["OFFSET_P"]

        # The OFFSET_T angle is in petal coordinates.  We express this in global
        # coordinates by transforming 2 points and computing the angle from that.

        prot = petal_rot[d["PETAL_ID"]]
        ptrans = petal_trans[d["PETAL_ID"]]
        offset_pt_rad = np.radians(d["OFFSET_T"])
        offset_pt_pnts = np.array(
            [
                [d["OFFSET_X_CS5"], d["OFFSET_Y_CS5"], 0.0],
                [
                    3.0 * np.cos(offset_pt_rad) + d["OFFSET_X_CS5"],
                    3.0 * np.sin(offset_pt_rad) + d["OFFSET_Y_CS5"],
                    0.0,
                ],
            ]
        )
        offset_fp_pnts = prot.dot(offset_pt_pnts.T).T + np.tile(ptrans, 2).reshape(
            (2, -1)
        )
        offset_fp_rad = np.arctan2(
            offset_fp_pnts[1, 1] - offset_fp_pnts[0, 1],
            offset_fp_pnts[1, 0] - offset_fp_pnts[0, 0],
        )
        offset_fp_deg = np.degrees(offset_fp_rad)
        if offset_fp_deg < -180.0:
            offset_fp_deg += 360.0
        if offset_fp_deg > 180.0:
            offset_fp_deg -= 360.0

        # # Sanity check that the alignment-transformed offset is "close" to the
        # # expected value.
        # petalrot_check = (float(7 + fp["PETAL"][r]) * 36.0) % 360.0
        # petalrot_check += d["OFFSET_T"]
        # if petalrot_check < -180.0:
        #     petalrot_check += 360.0
        # if petalrot_check > 180.0:
        #     petalrot_check -= 360.0
        # print(
        #     "device {}, petal {}, offset_t = {}, check = {}".format(
        #         fp["DEVICE"][r], fp["PETAL"][r], offset_fp_deg, petalrot_check
        #     ),
        #     flush=True,
        # )

        fp["OFFSET_T"][r] = offset_fp_deg

        # Set the FIBER
        fp["FIBER"][r] = -1
        if fp["SLITBLOCK"][r] >= 0:
            fp["FIBER"][r] = (
                fp["PETAL"][r] * 500 + fp["SLITBLOCK"][r] * 25 + fp["BLOCKFIBER"][r]
            )

        # Handle the keepout polygons for this device.  Every row specifies a keepout
        # even if many are close to identical.  Here we build up a dictionary of unique
        # keepouts and find the one to use for this positioner.

        dktheta = np.transpose(np.array(ast.literal_eval(d["KEEPOUT_T"])))
        dpolytheta = dict()
        dpolytheta["circles"] = list()
        dpolytheta["segments"] = collision_to_segments(dktheta)

        dkphi = np.transpose(np.array(ast.literal_eval(d["KEEPOUT_P"])))
        dpolyphi = dict()
        dpolyphi["circles"] = list()
        dpolyphi["segments"] = collision_to_segments(dkphi)

        kp = dict(excl["default"])
        kp["theta"] = dpolytheta
        kp["phi"] = dpolyphi
        pname = hash_exclusion(kp)[:16]

        if pname not in excl:
            excl[pname] = kp

        state["EXCLUSION"][r] = pname

        # Now set rest of the state table

        state["TIME"][r] = state_time_str
        state["LOCATION"][r] = d["PETAL_LOC"] * 1000 + d["DEVICE_LOC"]

        # Starting point is "good" with nominal MIN_P
        state["STATE"][r] = valid_states["OK"]
        state["MIN_P"][r] = d["MIN_P"]

        # Even if we are simulating, we want to mark both broken fibers and
        # non-movable positioners.
        if not d["FIBER_INTACT"]:
            state["STATE"][r] |= valid_states["BROKEN"]
        if d["DEVICE_CLASSIFIED_NONFUNCTIONAL"]:
            state["STATE"][r] |= valid_states["STUCK"]

        if not sim:
            # We want to also check for retracted positioners
            if d["CLASSIFIED_AS_RETRACTED"]:
                # This positioner is retracted.  Set the exclusion to the retracted
                # one and also limit the phi angle range.
                state["STATE"][r] |= valid_states["RESTRICT"]
                state["EXCLUSION"][r] = "retracted"
                # The focalplane cal information defines the Eo_Phi angle to be the
                # minimum Phi angle relative to the coordinate axis, not the offset.
                # So to get MIN_P we must subtract the offset.
                state["MIN_P"][r] = eo_phi - d["OFFSET_P"]

        # The other positioner angles in the state are just the same as nominal
        state["MAX_P"][r] = d["MAX_P"]
        state["MIN_T"][r] = d["MIN_T"]
        state["MAX_T"][r] = d["MAX_T"]
        # If the device is not movable, track its current estimated location.
        if state["STATE"][r] & stuck_or_broken:
            state["POS_P"][r] = d["POS_P"]
            state["POS_T"][r] = d["POS_T"]
        else:
            state["POS_P"][r] = 0.0
            state["POS_T"][r] = 0.0

    return (fp, state, excl, state_time_str)

def create_from_calibs(
    calib_file,
    out_dir=None,
    reset=False,
    sim_good=False,
    commit=False,
    test=False,
    fibermaps=None,
):
    """Construct a DESI focalplane from a calibration dump.

    This uses a dump from the online system and compares it to the current
    latest focalplane model and state.

    Args:
        calib_file (str):  Path to the calibration dump.
        out_dir (str):  Override the output directory for testing.  Default writes to
            $DESIMODEL/data/focalplane/
        reset (bool):  If True, ignore the current focalplane model and
            create a new model from this calibration dump.  Default compares
            the new focalplane to the old and looks for changes in device
            state.  These changes are appended to the existing log.
        sim_good (bool):  If True, clear all transient state issues and set hardware
            to be as "good as possible", for use in simulations.
        commit (bool):  If True, attempt to commit the result.
        test (bool):  If True, perform all operations but do not update any files.
        fibermaps (list):  Override list of tuples (DocDB number,
            DocDB version, DocDB csv file) of where to find the petal mapping
            files.

    Returns:
        None

    """
    log = get_logger()

    if out_dir is None:
        out_dir = os.path.join(datadir(), "focalplane")
        test_svn = os.path.join(datadir(), ".svn")
        if not os.path.isdir(test_svn):
            test_svn = os.path.join(os.path.dirname(datadir()), ".svn")
            if not os.path.isdir(test_svn):
                msg = "Output data directory:  {}".format(out_dir)
                msg += "\nis not inside an svn checkout.  You will not be able to"
                msg += "\ncommit these changes without copying them to a checkout."
                log.warning(msg)
    else:
        log.warning("Using debug output directory %s", out_dir)
        log.warning("Files cannot be used until placed in $DESIMODEL/data/focalplane")

    # Get the model from the calib file
    log.info("Loading calibration dump from %s ...", calib_file)
    fpcal = load_fp_calibs(calib_file)

    log.info("Converting calibration format ...")
    fp, state, excl, date_str = convert_fp_calibs(fpcal, sim=sim_good)

    log.info("Calibration data retrieval date = %s", date_str)

    if reset:
        # Ignore any previous focalplane info and dump out what we have

        log.info("Writing new focalplane model- ignoring previous ones...")
        out_fp_file = os.path.join(out_dir, "desi-focalplane_{}.ecsv".format(date_str))
        out_excl_file = os.path.join(
            out_dir, "desi-exclusion_{}.json.gz".format(date_str)
        )
        out_state_file = os.path.join(out_dir, "desi-state_{}.ecsv".format(date_str))

        log.info("  Output focalplane:  %s", out_fp_file)
        log.info("  Output state     :  %s", out_state_file)
        log.info("  Output exclusions:  %s", out_excl_file)

        if test:
            log.info("Running with test==True, skipping file writes.")
            return

        fp.write(out_fp_file, format="ascii.ecsv", overwrite=True)

        state.write(out_state_file, format="ascii.ecsv", overwrite=True)

        with gzip.open(out_excl_file, "wt", encoding="utf8") as pf:
            json.dump(excl, pf, indent=4)

        if commit:
            cmesg = "Creating new focalplane model from DB sync {}".format(date_str)
            try:
                sp.check_call(["svn", "update"], cwd=out_dir)
                sp.check_call(
                    [
                        "svn",
                        "add",
                        out_fp_file,
                        out_excl_file,
                        out_state_file,
                    ],
                    cwd=out_dir,
                )
                sp.check_call(["svn", "commit", "-m", cmesg], cwd=out_dir)
                sp.check_call(["svn", "update"], cwd=out_dir)
            except sp.CalledProcessError:
                log.error("svn update / commit returned an error")
    else:
        # Load the current focalplane and just update the state

        # Get the focalplane from the current datestamp
        cur_date = datetime.datetime.strptime(date_str, "%Y-%m-%dT%H:%M:%S%z")
        oldtime = cur_date

        oldfp, oldexcl, oldstate, oldtmstr = load_focalplane(oldtime)

        log.info("Comparing generated focalplane to one from %s", oldtmstr)

        # Compare the old and new.
        checkcols = set(fp.colnames)
        diff = device_compare(oldfp, fp, list(checkcols))

        if len(diff) > 0:
            device_printdiff(diff)
            msg = (
                "Existing focalplane device properties have changed."
                "  Use the 'reset' option to start with a new focalplane model."
            )
            log.error(msg)
            return

        # We got this far, which means that the focalplane data agrees.  Now
        # Look for differences in the state.
        checkcols = set(state.colnames)
        checkcols.remove("TIME")
        state_diff = device_compare(oldstate, state, list(checkcols))

        if len(state_diff) == 0:
            log.info("New focalplane state is identical, no action needed.")
            return

        # We must have some changes, load the full table and append.
        state_file = os.path.join(out_dir, "desi-state_{}.ecsv".format(oldtmstr))
        st = Table.read(state_file, format="ascii.ecsv")

        # If needed, promote the "TIME" state column to the new column type
        if len(st["TIME"][0]) < 21:
            newt = Column(
                name="TIME",
                length=len(st["TIME"]),
                dtype=np.dtype("a30"),
                description="The timestamp of the event (UTC, ISO format)",
            )
            for row, oldt in enumerate(st["TIME"]):
                nt = datetime.datetime.strptime(oldt, "%Y-%m-%dT%H:%M:%S")
                nt = nt.replace(tzinfo=datetime.timezone.utc)
                newt[row] = nt.isoformat()
            st.replace_column("TIME", newt)

        excl_file = None
        for test_ext in ["json.gz", "yaml.gz"]:
            excl_file = os.path.join(
                out_dir,
                "desi-exclusion_{}.{}".format(oldtmstr, test_ext)
            )
            if os.path.isfile(excl_file):
                break
            else:
                excl_file = None
        if excl_file is None:
            msg = "Attempting to sync to an ancient focalplane model with uncompressed YAML exclusion polygons"
            raise RuntimeError(msg)

        tmp_state = "{}.tmp".format(state_file)
        prev_state = "{}.previous".format(state_file)
        tmp_excl = "{}.tmp".format(excl_file)
        prev_excl = "{}.previous".format(excl_file)

        device_printdiff(state_diff)

        n_new_states = 0
        n_new_excl = 0
        for loc, df in state_diff.items():
            if df["old"] is None or df["new"] is None:
                # This should never happen, since it means that the LOCATION
                # value does not exist in either the previous or current
                # state.  We already checked for that above.
                msg = "LOCATION {} missing from old or new state.  Should never happen!".format(
                    loc
                )
                raise RuntimeError(msg)

            row = [df["new"][col] for col in df["new"].dtype.names]
            st.add_row(row)
            n_new_states += 1

            new_excl = str(
                df["new"]["EXCLUSION"].tobytes().rstrip(b"\x00"), encoding="utf-8"
            )
            if new_excl not in oldexcl:
                oldexcl[new_excl] = excl[new_excl]
                n_new_excl += 1

        log.info("Updating focalplane:  %s", oldtmstr)
        log.info("  State log appending %d rows", n_new_states)
        if n_new_excl > 0:
            log.info("  Adding %d new exclusion shapes", n_new_excl)
        else:
            log.info("  No new exclusion shapes, not updating file")

        if test:
            log.info("Running with test==True, skipping file writes.")
            return

        # Write to temp file then move into place
        st.write(tmp_state, format="ascii.ecsv", overwrite=True)
        shutil.copy2(state_file, prev_state)
        os.rename(tmp_state, state_file)

        # If we updated any exclusions, write a new file
        if n_new_excl > 0:
            with gzip.open(tmp_excl, "wt", encoding='utf8') as pf:
                if re.match(r".*json.*", excl_file) is not None:
                    # New format
                    json.dump(
                        oldexcl,
                        pf,
                        indent=4
                    )
                else:
                    # Old format
                    yaml.dump(
                        oldexcl,
                        stream=pf,
                        encoding="utf-8",
                        default_flow_style=False
                    )

            shutil.copy2(excl_file, prev_excl)
            os.rename(tmp_excl, excl_file)

        if commit:
            cmesg = "Appending DB sync {} to focalplane model {}".format(
                date_str, oldtmstr
            )
            try:
                sp.check_call(["svn", "update"], cwd=out_dir)
                sp.check_call(
                    [
                        "svn",
                        "commit",
                        "-m",
                        cmesg,
                    ],
                    cwd=out_dir,
                )
                sp.check_call(["svn", "update"], cwd=out_dir)
            except sp.CalledProcessError:
                log.error("svn update / commit returned an error")


"""`AlphaIMS`, `AlphaAMS`"""
import numpy as np
from collections import OrderedDict

from .base import ProsthesisSystem
from .electrodes import SquareElectrode, DiskElectrode
from .electrode_arrays import ElectrodeGrid

class AlphaIMS(ProsthesisSystem):
    """Alpha-IMS

    This class creates an Alpha-IMS array with 1500 photovoltaic pixels (each
    50um in diameter) as described in [Stingl2013]_, and places it in the
    subretinal space, such that the center of the array is located at (x,y,z),
    given in microns, and the array is rotated by rotation angle ``rot``,
    given in degrees.

    The device consists of 1500 50um-wide square pixels, arranged on a 39x39
    rectangular grid with 72um pixel pitch.

    The array is oriented upright in the visual field, such that an
    array with center (0,0) has the top three rows lie in the lower
    retina (upper visual field).

    An electrode can be addressed by name, row/column index, or integer index
    (into the flattened array).

    .. note::

        Column order is reversed in a left-eye implant.

    Parameters
    ----------
    x/y/z : double
        3D location of the center of the electrode array.
        The coordinate system is centered over the fovea.
        Positive ``x`` values move the electrode into the nasal retina.
        Positive ``y`` values move the electrode into the superior retina.
        Positive ``z`` values move the electrode away from the retina into the
        vitreous humor (sometimes called electrode-retina distance).
        ``z`` can either be a list with 1500 entries or a scalar that is applied
        to all electrodes.
    rot : float
        Rotation angle of the array (deg). Positive values denote
        counter-clock-wise (CCW) rotations in the retinal coordinate
        system.
    eye : {'RE', 'LE'}, optional
        Eye in which array is implanted.
    preprocess : bool or callable, optional
        Either True/False to indicate whether to execute the implant's default
        preprocessing method whenever a new stimulus is assigned, or a custom
        function (callable).
    safe_mode : bool, optional
        If safe mode is enabled, only charge-balanced stimuli are allowed.

    Examples
    --------
    Create an Alpha-IMS array centered on the fovea, at 100um distance from
    the retina, rotated counter-clockwise by 5 degrees:

    >>> from pulse2percept.implants import AlphaIMS
    >>> AlphaIMS(x=0, y=0, z=100, rot=5)  # doctest: +NORMALIZE_WHITESPACE
    AlphaIMS(earray=ElectrodeGrid, eye='RE', preprocess=True,
             safe_mode=False, shape=(39, 39), stim=None)

    Get access to the third electrode in the top row (by name or by row/column
    index):

    >>> alpha_ims = AlphaIMS(x=0, y=0, z=100, rot=0)
    >>> alpha_ims['A3']  # doctest: +NORMALIZE_WHITESPACE
    SquareElectrode(a=50.0, activated=True, name='A3',
                    x=-1224.0, y=-1368.0, z=100.0)
    >>> alpha_ims[0, 2]  # doctest: +NORMALIZE_WHITESPACE
    SquareElectrode(a=50.0, activated=True, name='A3',
                    x=-1224.0, y=-1368.0, z=100.0)

    """
    # Frozen class: User cannot add more class attributes
    __slots__ = ('shape',)

    def __init__(self, x=0, y=0, z=-100, rot=0, eye='RE', stim=None,
                 preprocess=True, safe_mode=False):
        self.eye = eye
        self.preprocess = preprocess
        self.safe_mode = safe_mode
        self.shape = (39, 39)
        elec_width = 50.0  # um
        e_spacing = 72.0  # um

        # The user might provide a list of z values for each of the
        # 378 resulting electrodes, not for the 22x19 initial ones.
        # In this case, don't pass it to ElectrodeGrid, but overwrite
        # the z values later:
        overwrite_z = isinstance(z, (list, np.ndarray))
        zarr = -100.0 if overwrite_z else z
        self.earray = ElectrodeGrid(self.shape, e_spacing, x=x, y=y, z=zarr,
                                    rot=rot, etype=SquareElectrode,
                                    a=elec_width)

        # Unfortunately, in the left eye the labeling of columns is reversed...
        if eye == 'LE':
            # FIXME: Would be better to have more flexibility in the naming
            # convention. This is a quick-and-dirty fix:
            names = self.earray.electrode_names
            objects = self.earray.electrode_objects
            names = np.array(names).reshape(self.earray.shape)
            # Reverse column names:
            for row in range(self.earray.shape[0]):
                names[row] = names[row][::-1]
            # Build a new ordered dict:
            electrodes = OrderedDict([])
            for name, obj in zip(names.ravel(), objects):
                electrodes.update({name: obj})
            # Assign the new ordered dict to earray:
            self.earray._electrodes = electrodes

        # Remove electrodes:
        extra_elecs = ['AM39', 'AL39', 'AK39', 'AJ39', 'AI39', 'AH39', 'AG39',
                       'AF39', 'AE39', 'AD39', 'AC39',
                       'AM38', 'AL38', 'AK38', 'AJ38', 'AI38', 'AH38', 'AG38',
                       'AF38', 'AE38', 'AD38']
        for elec in extra_elecs:
            self.earray.remove_electrode(elec)

        # Now that the superfluous electrodes have been deleted, adjust the
        # z values:
        if overwrite_z:
            # Specify different height for every electrode in a list:
            z_arr = np.asarray(z).flatten()
            if z_arr.size != self.n_electrodes:
                raise ValueError("If `z` is a list, it must have %d entries, "
                                 "not %d." % (self.n_electrodes, z_arr.size))
            for elec, z_elec in zip(self.earray.electrode_objects, z):
                elec.z = z_elec

        # Beware of race condition: Stim must be set last, because it requires
        # indexing into self.electrodes:
        self.stim = stim

    def _pprint_params(self):
        """Return dict of class attributes to pretty-print"""
        params = super()._pprint_params()
        params.update({'shape': self.shape})
        return params

class AlphaAMS(ProsthesisSystem):
    """Alpha-AMS

    This class creates an Alpha-AMS array with 1600 photovoltaic pixels (each
    30um in diameter) as described in [Stingl2017]_, and places it in the
    subretinal space, such that the center of the array is located at (x,y,z),
    given in microns, and the array is rotated by rotation angle ``rot``,
    given in degrees.

    The device consists of 1600 30um-wide round pixels, arranged on a 40x40
    rectangular grid with 70um pixel pitch.

    The array is oriented upright in the visual field, such that an
    array with center (0,0) has the top three rows lie in the lower
    retina (upper visual field), as shown below:

    An electrode can be addressed by name, row/column index, or integer index
    (into the flattened array).

    .. note::

        Column order is reversed in a left-eye implant.

    Parameters
    ----------
    x/y/z : double
        3D location of the center of the electrode array.
        The coordinate system is centered over the fovea.
        Positive ``x`` values move the electrode into the nasal retina.
        Positive ``y`` values move the electrode into the superior retina.
        Positive ``z`` values move the electrode away from the retina into the
        vitreous humor (sometimes called electrode-retina distance).
        ``z`` can either be a list with 1600 entries or a scalar that is applied
        to all electrodes.
    rot : float
        Rotation angle of the array (deg). Positive values denote
        counter-clock-wise (CCW) rotations in the retinal coordinate
        system.
    eye : {'RE', 'LE'}, optional
        Eye in which array is implanted.
    preprocess : bool or callable, optional
        Either True/False to indicate whether to execute the implant's default
        preprocessing method whenever a new stimulus is assigned, or a custom
        function (callable).
    safe_mode : bool, optional
        If safe mode is enabled, only charge-balanced stimuli are allowed.

    Examples
    --------
    Create an AlphaAMS array centered on the fovea, at 100um distance from
    the retina, rotated counter-clockwise by 5 degrees:

    >>> from pulse2percept.implants import AlphaAMS
    >>> AlphaAMS(x=0, y=0, z=100, rot=5)  # doctest: +NORMALIZE_WHITESPACE
    AlphaAMS(earray=ElectrodeGrid, eye='RE', preprocess=True,
             safe_mode=False, shape=(40, 40), stim=None)

    Get access to the third electrode in the top row (by name or by row/column
    index):

    >>> alpha_ims = AlphaAMS(x=0, y=0, z=100, rot=0)
    >>> alpha_ims['A3']  # doctest: +NORMALIZE_WHITESPACE
    DiskElectrode(activated=True, name='A3', r=15.0, x=-1225.0,
                  y=-1365.0, z=100.0)
    >>> alpha_ims[0, 2]  # doctest: +NORMALIZE_WHITESPACE
    DiskElectrode(activated=True, name='A3', r=15.0, x=-1225.0,
                  y=-1365.0, z=100.0)

    """
    # Frozen class: User cannot add more class attributes
    __slots__ = ('shape',)

    def __init__(self, x=0, y=0, z=0, rot=0, eye='RE', stim=None,
                 preprocess=True, safe_mode=False):
        self.eye = eye
        self.preprocess = preprocess
        self.safe_mode = safe_mode
        self.shape = (40, 40)
        elec_radius = 15.0
        e_spacing = 70.0  # um

        self.earray = ElectrodeGrid(self.shape, e_spacing, x=x, y=y, z=z,
                                    rot=rot, etype=DiskElectrode,
                                    r=elec_radius)

        # Beware of race condition: Stim must be set last, because it requires
        # indexing into self.electrodes:
        self.stim = stim

        # Set left/right eye:
        # Unfortunately, in the left eye the labeling of columns is reversed...
        if eye == 'LE':
            # FIXME: Would be better to have more flexibility in the naming
            # convention. This is a quick-and-dirty fix:
            names = self.earray.electrode_names
            objects = self.earray.electrode_objects
            names = np.array(names).reshape(self.earray.shape)
            # Reverse column names:
            for row in range(self.earray.shape[0]):
                names[row] = names[row][::-1]
            # Build a new ordered dict:
            electrodes = OrderedDict([])
            for name, obj in zip(names.ravel(), objects):
                electrodes.update({name: obj})
            # Assign the new ordered dict to earray:
            self.earray._electrodes = electrodes

    def _pprint_params(self):
        """Return dict of class attributes to pretty-print"""
        params = super()._pprint_params()
        params.update({'shape': self.shape})
        return params


#!/usr/bin/env python3
#
# (c) Copyright Rosetta Commons Member Institutions.
# (c) This file is part of the Rosetta software suite and is made available under license.
# (c) The Rosetta software is developed by the contributing members of the Rosetta Commons.
# (c) For more information, see http://www.rosettacommons.org. Questions about this can be
# (c) addressed to University of Washington CoMotion, email: license@uw.edu.
'''
Contains helper functions that defines Rosetta-style topology for given Molecule including:
- AtomTree, Chi angles, Ring, Bond order, and so on.

Author: Hahnbeom Park and Frank DiMaio
'''

import sys
import math
import numpy as np
import numpy.linalg
import scipy.sparse.csgraph
import scipy.linalg
from Types import *
from BasicClasses import AtomClass, RingClass
from AtomTypeClassifier import AtomTypeClassifier
from utils import dihedral

CUTOFF_PLANAR = 40.0 # GZ: this is just a guess

def setup(mol,option): # mol: Molecule type
    '''Main function setting up all the toplogy-related things of the input Molecule(=="mol") '''
    mol.vatms = []
    mol.angles = []
    mol.torsions = []
    mol.nbratom = -1
    mol.nbrradius = 0.0
    mol.atms_aro = []
    mol.atms_puckering = []
    mol.atms_ring = [] #including aromatic ring
    mol.ATorder = []
    mol.rings_strained = [] #3/4 membered
    mol.rings_aro = [] #only aromatic
    mol.rings_pucker = [] #puckering
    mol.rings_sugar = [] #subset of puckering rings
    mol.rings = [] #aromatic + puckering

    assign_bonds(mol)
    assign_angles(mol)
    assign_torsions(mol)
    detect_rings(mol,option)

    assign_hybridization_if_missing(mol) # should come after ring assignment
    setup_nbratom(mol) #should come after ring assignments
    define_icoord(mol) #should come after nbr setup

    classifier = AtomTypeClassifier()
    classifier.apply_to_molecule(mol)
    classifier.assert_H(mol)

    # should come after atom types classified
    define_conjugation(mol)
    define_rotable_torsions(mol)

# Functions from here are pieces of setup

def setup_nbratom(mol):
    '''Setup residue nbratom & nbrradius field'''
    com = np.sum(mol.xyz,axis=0) / len(mol.atms)
    dists = np.sum(np.square(mol.xyz - com[None,:]),axis=-1) # dist_squared
    maxdis = math.sqrt(max(dists))

    # 1) dont let puckering ring be root
    dists[np.in1d(np.arange(len(mol.atms)), mol.atms_puckering)] = 99999.0

    nbonds = [len(atm.bonds) for atm in mol.atms]
    for ring in mol.rings: #correct for cut_bonds
        (i,j) = ring.cut_bond
        if i == None: continue
        nbonds[i] -=1
        nbonds[j] -=1

    # 3) also exclude hydrogens
    # prioritize ring-atom with >2 bonds
    for i,atm in enumerate(mol.atms):
        if nbonds[i] <=1: dists[i] = 99999.0
        elif nbonds[i] ==2: dists[i] = 999.0
        if atm.is_H: dists[i] = 99999.0

    mol.nbratom = np.argmin(dists)
    mol.nbrradius = (maxdis+1.5)*2 #safe

def assign_bonds(mol):
    for bond in mol.bonds:
        is_H1 = (mol.atms[bond.atm1].atype==2)
        is_H2 = (mol.atms[bond.atm2].atype==2)
        mol.atms[bond.atm1].add_bond(bond.atm2,bond.order)
        mol.atms[bond.atm2].add_bond(bond.atm1,bond.order)

        if is_H1:
            mol.atms[bond.atm2].has_H = True
        if is_H2:
            mol.atms[bond.atm1].has_H = True

        if mol.atms[bond.atm1].atype == 1 and mol.atms[bond.atm2].atype in POLAR_ATOMS:
            mol.atms[bond.atm1].connected_to_polar = True
        if mol.atms[bond.atm2].atype == 1 and mol.atms[bond.atm1].atype in POLAR_ATOMS:
            mol.atms[bond.atm2].connected_to_polar = True

    # sanity check if any atm is not connected to anything
    atms_unconnected = []
    for atm in mol.atms:
        if len(atm.bonds) == 0:
            atms_unconnected.append(atm.name)
    if atms_unconnected != []:
        sys.exit('ERROR: Atom found not connected to any other atom:'+' '.join(atms_unconnected))

# Currently used only for 'nh' case: See "0" in SPECIAL_HYBRIDS to see which atypes in mol2 comes here
def assign_hybridization_if_missing(mol):
    n_to_assign = 0
    to_assign = []
    for iatm,atm in enumerate(mol.atms):
        if ATYPES[atm.atype] in ATYPES_HYBRID:
            n_to_assign += 1
            if atm.hyb == 0:
                to_assign.append(iatm)

    for iatm in to_assign:
        atm = mol.atms[iatm]
        attached_to_aro = False
        attached_to_nonaro = False
        nH_attached = 0
        for jatm,order in atm.bonds:
            if jatm in mol.atms_aro:
                attached_to_aro = True
            elif mol.atms[jatm].is_H:
                nH_attached += 1
            else:
                attached_to_nonaro = True

        if atm.atype == ATYPES.index('N') and attached_to_aro:
            if nH_attached == 2: #aro-NH2
                atm.hyb = 2
            elif not attached_to_nonaro:
                atm.hyb = 2
            else:
                atm.hyb = 3
        else:
            if len(atm.bonds) < 3:
                atm.hyb = 2
            else:
                #measure improper torsion
                neigh = [j for j,order in atm.bonds][:3]
                imp = dihedral(mol.xyz[neigh[0]],mol.xyz[neigh[1]],mol.xyz[neigh[2]],mol.xyz[iatm])
                if abs(imp) < 10.0: #dev less than 20 degree
                    atm.hyb = 2
                else:
                    atm.hyb = 3

        # BELOW FAILS for GAFFtype nh!
        # Use user input structure instead above (there is no other easy way)
        #bondorders = [bond[1] for bond in atm.bonds]
        #maxorder = max(bondorders)
        #if 1 in bondorders:
        #    atm.hyb = 3
        #elif maxorder == 9:
        #    atm.hyb = 9
        #elif maxorder == 2:
        #    atm.hyb = 2
        #elif maxorder == 3:
        #    atm.hyb = 1
        #else:
        #    print('WARNING: Unknown atype in input mol2 for %s; assigning hyb=2'%(atm.name))
        print('Hybridization state for %d:%s assigned as %s'%(iatm,atm.name,atm.hyb))

def assign_angles(mol):
    '''Assigns all bond-angles (not only in AtomTree)'''
    for i,bond1 in enumerate(mol.bonds[:-1]):
        for bond2 in mol.bonds[i+1:]:
            if bond1.atm1 == bond2.atm1:
                angle = (bond1.atm2,bond1.atm1,bond2.atm2)
            elif bond1.atm1 == bond2.atm2:
                angle = (bond1.atm2,bond1.atm1,bond2.atm1)
            elif bond1.atm2 == bond2.atm1:
                angle = (bond1.atm1,bond1.atm2,bond2.atm2)
            elif bond1.atm2 == bond2.atm2:
                angle = (bond1.atm1,bond1.atm2,bond2.atm1)
            else:
                continue
            mol.angles.append(angle)

def assign_torsions(mol):
    '''Assigns all torsion angles (not only in AtomTree)'''
    bonds = [(bond.atm1,bond.atm2) for bond in mol.bonds]

    torsions_heavy = []
    torsions_H = []
    orders_heavy = []
    orders_H = []
    if mol.option.verbose:
        print('bonds: ', [(bond.atm1,bond.atm2) for bond in mol.bonds])
    for i,bond1 in enumerate(mol.bonds[:-1]):
        for bond2 in mol.bonds[i+1:]:
            if (bond1.atm1 == bond2.atm1 or bond1.atm1 == bond2.atm2 or \
                bond1.atm2 == bond2.atm1 or bond1.atm2 == bond2.atm2): continue

            if (bond1.atm1,bond2.atm1) in bonds:
                a,b,c,d = (bond1.atm2,bond1.atm1,bond2.atm1,bond2.atm2)
            elif (bond1.atm1,bond2.atm2) in bonds:
                a,b,c,d = (bond1.atm2,bond1.atm1,bond2.atm2,bond2.atm1)
            elif (bond1.atm2,bond2.atm1) in bonds:
                a,b,c,d = (bond1.atm1,bond1.atm2,bond2.atm1,bond2.atm2)
            elif (bond1.atm2,bond2.atm2) in bonds:
                a,b,c,d = (bond1.atm1,bond1.atm2,bond2.atm2,bond2.atm1)
            elif (bond2.atm1,bond1.atm1) in bonds:
                a,b,c,d = (bond2.atm2,bond2.atm1,bond1.atm1,bond1.atm2)
            elif (bond2.atm2,bond1.atm1) in bonds:
                a,b,c,d = (bond2.atm1,bond2.atm2,bond1.atm1,bond1.atm2)
            elif (bond2.atm1,bond1.atm2) in bonds:
                a,b,c,d = (bond2.atm2,bond2.atm1,bond1.atm2,bond1.atm1)
            elif (bond2.atm2,bond1.atm2) in bonds:
                a,b,c,d = (bond2.atm1,bond2.atm2,bond1.atm2,bond1.atm1)
            else:
                continue

            ibond = bonds.index((b,c))
            ## bug fix Sep25 2018: if having only hydrogens at either end
            if mol.atms[a].is_H or mol.atms[d].is_H:
                torsions_H.append((a,b,c,d))
                orders_H.append(mol.bonds[ibond].order)
            else:
                torsions_heavy.append((a,b,c,d))
                orders_heavy.append(mol.bonds[ibond].order)

    # sort in heavy-atom torsions followed by H-only torsions
    #mol.torsion_orderso = orders_heavy + orders_H #unused
    mol.torsions = torsions_heavy + torsions_H

    if mol.option.verbose:
        print('Torsions: ', mol.torsions)

def search_special_biaryl_ring(mol,ring):
    '''Add extra rotatable bonds between ring & nonring functional groups
    (by going through hard-coded cases)'''

    biaryl_pivot_extra = []
    # [(A,B),..] : ring=A becomes special biaryl if any of ring=A-B connection is detected
    # Nin definition changed as of Oct242018; All the non-aromatic 1H sp2N goes to NG21
    # non-biaryl (== conjugated): NG21=O*
    special_biaryl_to_ring = [(ACLASS_ID.index('CDp'),ACLASS_ID.index('OG2')),
                              (ACLASS_ID.index('CDp'),ACLASS_ID.index('Oal')),
                              (ACLASS_ID.index('CDp'),ACLASS_ID.index('Oad')),

                              # ring-N=C; (RingN-C)-NG21 are excluded outside here
                              (ACLASS_ID.index('NG21'),ACLASS_ID.index('CR')),
                              (ACLASS_ID.index('NG21'),ACLASS_ID.index('CRp')),
                              (ACLASS_ID.index('NG21'),ACLASS_ID.index('CD')),
                              (ACLASS_ID.index('NG21'),ACLASS_ID.index('CDp')),
                              (ACLASS_ID.index('NG21'),ACLASS_ID.index('CSp')),
                              (ACLASS_ID.index('NG21'),ACLASS_ID.index('CS')),
                              (ACLASS_ID.index('NG21'),ACLASS_ID.index('CS1')),
                              (ACLASS_ID.index('NG21'),ACLASS_ID.index('CS2')),

                              # ring-N-SOn
                              #(ACLASS_ID.index('Nin'),ACLASS_ID.index('SG5')),
                              (ACLASS_ID.index('NG21'),ACLASS_ID.index('SG5')),

                              # ring-amide
                              (ACLASS_ID.index('Nad'),ACLASS_ID.index('CDp')), # treat energy separately

                              # ring-C=N*
                              (ACLASS_ID.index('CDp'),ACLASS_ID.index('Nad')),
                              (ACLASS_ID.index('CDp'),ACLASS_ID.index('Nam2')),
                              (ACLASS_ID.index('CDp'),ACLASS_ID.index('NG2')),
                              (ACLASS_ID.index('CDp'),ACLASS_ID.index('Nin')),

                              # R-CD1-CD1-
                              (ACLASS_ID.index('CD1'),ACLASS_ID.index('CD1')),
                              (ACLASS_ID.index('CD1'),ACLASS_ID.index('CD')),
                              (ACLASS_ID.index('CD'),ACLASS_ID.index('CD')),
                              (ACLASS_ID.index('CD'),ACLASS_ID.index('CD1')),

                              # R-CDx-R
                              (ACLASS_ID.index('CD'),ACLASS_ID.index('CR')),
                              (ACLASS_ID.index('CD1'),ACLASS_ID.index('CR')),
                              (ACLASS_ID.index('CDp'),ACLASS_ID.index('CR')),

                              # ring-C=halogen
                              (ACLASS_ID.index('CD'),ACLASS_ID.index('F')),
                              (ACLASS_ID.index('CD'),ACLASS_ID.index('Cl')),
                              (ACLASS_ID.index('CD'),ACLASS_ID.index('Br')),
                              (ACLASS_ID.index('CD'),ACLASS_ID.index('I')),
                              ]

    for a1 in ring.atms:
        a2_is_special_biaryl = False
        for a2,order in mol.atms[a1].bonds:
            a2class = mol.atms[a2].aclass
            # avoid from ring itself
            #if a2 in mol.atms_aro or a2 in mol.atms_puckering: continue
            if mol.in_same_ring(a1,a2): continue

            # Oct17: skip if not defined as single bond as input!
            if order > 1: continue

            further_connected = False
            for a3,order in mol.atms[a2].bonds:
                if a3 == a1: continue
                a3class = mol.atms[a3].aclass

                if (a2class,a3class) in special_biaryl_to_ring:
                    a2_is_special_biaryl = True
                    #break
                if len(mol.atms[a3].bonds) > 1:
                    further_connected = True

            if a2_is_special_biaryl and further_connected:
                #Final check using the input conformation
                #if the torsion around a1-a2 is planar, treat it as conjuated.
                #e.g UZARAP
                #GZ, Aug,20,2020
                if mol.option.opt.infer_conjugation and is_planar(mol, a1, a2):
                    continue
                biaryl_pivot_extra.append((a1,a2))

    return biaryl_pivot_extra

def check_duplication(rings,ring_i):
    for ring_j in rings:
        n = 0
        for j in ring_j.atms:
            if j in ring_i: n+=1
        if n == ring_j.natms:
            return ring_j.atms
    return False

def get_path(bond_tree,edge,alt_cut=[]):
    '''helper function in ring detection'''
    i,j = edge
    curr = j
    ring_i = [curr]
    bond_tree[i,j] = bond_tree[j,i] = 0
    if alt_cut != []: bond_tree[alt_cut[0],alt_cut[1]] = bond_tree[alt_cut[1],alt_cut[0]] = 0
    ds, preds = scipy.sparse.csgraph.shortest_path(bond_tree,return_predecessors=True)
    bond_tree[i,j] = bond_tree[j,i] = 1
    if alt_cut != []: bond_tree[alt_cut[0],alt_cut[1]] = bond_tree[alt_cut[1],alt_cut[0]] = 1

    if ds[i,j] > 9999: return False

    while (curr != i):
        curr = preds[i,curr]
        ring_i.append(curr)
    return ring_i

def detect_rings(mol,option):
    '''Detects ring & ring types (i.e. aro vs nonaro)'''
    mol.rings = []
    mol.atms_aro = []
    mol.atms_puckering = []

    bond_tree = np.zeros((len(mol.atms),len(mol.atms)))
    for bond in mol.bonds:
        bond_tree[bond.atm1,bond.atm2] = bond_tree[bond.atm2,bond.atm1] = 1

    mol.tree = scipy.sparse.csgraph.minimum_spanning_tree(bond_tree)
    treesymm = np.maximum(mol.tree.toarray(), mol.tree.T.toarray())
    cycle_edges = np.transpose(np.nonzero(np.triu(bond_tree != treesymm)))

    for i,j in cycle_edges:
        ring_i = get_path(bond_tree,(i,j))

        # this sometimes happens when cycle_edge connects two rings
        dupl_ring = check_duplication(mol.rings,ring_i)

        if dupl_ring: #search alternative path by cutting an edge of duplicating ring
            ring_i = get_path(bond_tree,(i,j),alt_cut=(dupl_ring[0],dupl_ring[-1]))
            if ring_i:
                dupl_ring = check_duplication(mol.rings,ring_i)
                if not dupl_ring:
                    mol.rings.append(RingClass(ring_i,cut=(ring_i[-1],ring_i[0])))

        else: #One cycle edge can form alternative ring
            mol.rings.append(RingClass(ring_i,cut=(ring_i[-1],ring_i[0])))
            ring_ii = get_path(bond_tree,(i,j),alt_cut = (ring_i[0],ring_i[1]))
            if ring_ii:
                dupl_ring = check_duplication(mol.rings,ring_ii)
                if not dupl_ring:
                    mol.rings.append(RingClass(ring_ii,cut=(ring_ii[-1],ring_ii[0])))

    for ring in mol.rings:
        classify_ring_type(mol,ring,option)

    mol.atms_ring = mol.atms_aro + mol.atms_puckering

def classify_ring_type(mol,ring,option):
    '''Classifies ring type'''
    is_aro = True

    #for long ring take into account of input geometry
    if ring.natms > 6:
        ring_xs = [mol.xyz[i] for i in ring.atms]
        ring_xs -= np.mean(ring_xs, axis = 0)
        cov = np.cov(ring_xs, rowvar = False)
        evals,_ = scipy.linalg.eigh(cov)
        nonplanarity = np.min(evals)
        if (nonplanarity>1e-2):  # be a bit permissive
            is_aro = False
            #print("nonplanarity? ", nonplanarity)
    else: #3~6 membered rings
        nsp2N = 0
        nsp3OS = 0
        nsp3 = 0
        for atm in ring.atms:
            if mol.atms[atm].hyb == 3:
                nsp3 += 1
                for j,order in mol.atms[atm].bonds:
                    if order not in [2,4]:
                        is_aro = False
                if ATYPES[mol.atms[atm].atype] in ['O','S']:
                    nsp3OS += 1
            elif mol.atms[atm].hyb == 2 and ATYPES[mol.atms[atm].atype] == 'N':
                nsp2N += 1

            # allow exception: 2-nitrogen 5-membered ring or one sp3 O/S ring
            if len(ring.atms) == 5:
                if (nsp3 == nsp3OS) or (nsp2N == 2):
                    is_aro = True
                else:
                    is_aro = False

    if ring.natms <= 4:
        ring.type = 2
        mol.rings_strained.append(ring)
        for atm in ring.atms:
            if atm not in mol.atms_strained:
                mol.atms_strained.append(atm)

    elif is_aro:
        ring.type = 4
        mol.rings_aro.append(ring)
        for atm in ring.atms:
            if atm not in mol.atms_aro:
                mol.atms_aro.append(atm)

    else: #otherwise puckering
        ring.type = 2 #non-ringsampling
        natms = len(ring.atms)
        has_nonsp3 = False
        for atm in ring.atms:
            if mol.atms[atm].hyb != 3:
                has_nonsp3 = True
                break

        if option.opt.report_puckering_chi:
            if (natms > option.opt.longest_puckering_ring):
                ring.type = 3 #long
            elif (natms > 4) and (not option.opt.ring_sampling_sp3_only or not has_nonsp3 ):
                ring.type = 1 #ring-sampling

        #sugar
        if len(ring.atms) == 5:
            n_sp3C = 0
            n_O = 0
            for iatm in ring.atms:
                atm = mol.atms[iatm]
                if ATYPES[atm.atype] == 'C' and atm.hyb == 3:
                    n_sp3C += 1
                if ATYPES[atm.atype] == 'O': n_O += 1
            if n_sp3C == 4 and n_O == 1:
                #ring.type = 5
                print('Sugar ring %s: '%mol.mol2file+'  %3s'*5%tuple([mol.atms[iatm].name for iatm in ring.atms]))
        mol.rings_pucker.append(ring)

        # puckering also addes into rings (for the biaryl assignment...)
        for atm in ring.atms:
            if atm not in mol.atms_puckering:
                mol.atms_puckering.append(atm)

def is_colinear(crd1, crd2, crd3, eps=1.0e-5):
    v1 = np.array(crd1) - np.array(crd2)
    v2 = np.array(crd3) - np.array(crd2)
    v4_0 = np.cross(v1,v2)
    if np.inner(v4_0,v4_0) < eps:
        return True
    return False

# determine if the torsion defined around the central atm1-atm2 bond is planar
def is_planar(mol, atm1, atm2):
    dihe = None
    # If one of the central atom is sp1, the torsion is not well defined around atm1-atm2
    # so we don't consider it as planar at this point.
    # GZ, Aug,20,2020
    if mol.atms[atm1].hyb == 1 or mol.atms[atm2].hyb == 1:
        return False
    for itor,(a1,a2,a3,a4) in enumerate(mol.torsions):
        if ( atm1 == a2 and atm2 == a3 ) or \
            ( atm1 == a3 and atm2 == a2 ):
            xyz0 = mol.xyz[a1]
            xyz1 = mol.xyz[a2]
            xyz2 = mol.xyz[a3]
            xyz3 = mol.xyz[a4]
            dihe = dihedral(xyz0,xyz1,xyz2,xyz3)
            break
    if dihe is None:
        print("Warning: No torsion defined around the single bond between atm: %d, %d. Set dihe=0 "%(atm1, atm2))
        dihe = 0.0
    if abs(dihe) >= CUTOFF_PLANAR and abs(dihe) <= (180-CUTOFF_PLANAR) :
        return False
    return True

def validate_order(nodes, parents):
    new_nodes = []
    cur_parent=-999
    children_nodes = []
    for i in nodes:
        if cur_parent != parents[i]:
            new_nodes.extend( np.sort(children_nodes) )
            children_nodes = [i]
            cur_parent = parents[i]
        else:
            children_nodes.append(i)
    new_nodes.extend( np.sort(children_nodes) )

    return new_nodes

# Scipy-version
def define_icoord(mol):
    '''AtomTree setup using scipy graph construct'''
    nodes,parents = scipy.sparse.csgraph.breadth_first_order(mol.tree, mol.nbratom, directed=False)
    nodes = validate_order(nodes, parents)
    first_children = np.zeros_like(parents)
    mol.ATorder = nodes
    colinear_child = -9999
    colinear_parent = -9999
    for i in nodes:
        par_i = parents[i]
        gp_i = -9999 if (par_i==-9999) else parents[par_i]
        ggp_i = -9999 if (gp_i==-9999) else parents[gp_i]

        # 1: root corrections
        if (i == nodes[0] or i == nodes[1] or i == nodes[2]):
            par_i = nodes[0]
            gp_i = nodes[1]
            ggp_i = nodes[2]

        # 2: near-root corrections
        if (gp_i == -9999):
            # 3+ child of root
            gp_i = nodes[1]
            ggp_i = nodes[2]

        if (ggp_i == -9999):
            if (par_i == nodes[1]):
                ggp_i = nodes[2]
            else:
                ggp_i = nodes[1]

        if is_colinear(mol.xyz[par_i], mol.xyz[gp_i], mol.xyz[ggp_i]):
            if colinear_parent != par_i:
                colinear_parent = par_i
                colinear_child = i
            else:
                ggp_i = colinear_child
                colinear_child = i

        mol.atms[i].root = par_i
        mol.atms[i].groot = (gp_i,ggp_i)

        # ring virtuals
        if (i in mol.atms_puckering):
            for rnum,ring in enumerate(mol.rings):
                (k,l) = ring.cut_bond
                if not (k==i or l==i): continue
                if ring.type != 1: continue

                if (i==k):
                    vrt_i = l
                    vtag = "V%dl"%rnum
                else:
                    vrt_i = k
                    vtag = "V%du"%rnum

                # define virtual atoms
                atm = AtomClass(vtag,"X",0,0.0)
                atm.vrt_i = vrt_i
                atm.root = i
                atm.ring_index = rnum

                if (gp_i == i): # near-root corr
                    atm.groot = (par_i,nodes[2])
                else:
                    atm.groot = (par_i,gp_i)
                mol.vatms.append(atm)

def is_biaryl_ring(mol,ring1,ring2):
    '''Check conjugation between two rings'''
    # first check if shares any atom
    for atm1 in ring1.atms:
        for atm2 in ring2.atms:
            if atm1 == atm2:
                return False
    # Then figure out if any atm pair connected
    for atm1 in ring1.atms:
        for atm2 in ring2.atms:
            for bond in mol.bonds:
                if (atm1,atm2) != (bond.atm1,bond.atm2) and \
                   (atm1,atm2) != (bond.atm2,bond.atm1):
                    continue

                #check if atm1,atm2 are connected by ring
                is_connected_by_ring = False
                for ring in mol.rings:
                    if atm1 in ring.atms and atm2 in ring.atms:
                        is_connected_by_ring = True
                        break
                if is_connected_by_ring:
                    continue
                if mol.option.opt.infer_conjugation:
                    if ACLASS_ID[mol.atms[atm1].aclass] in CONJUGATING_ACLASSES and \
                        ACLASS_ID[mol.atms[atm2].aclass] in CONJUGATING_ACLASSES and \
                        is_planar(mol, atm1, atm2):
                        continue
                return (atm1,atm2)
    return False

def define_conjugation(mol):
    '''Defines conjugations in bonds in Molecule based on heuristics'''
    # 1. Rotable ring-ring
    mol.biaryl_rings = []
    mol.biaryl_pivots = []
    if len(mol.rings) > 1:
        for i,ring1 in enumerate(mol.rings[:-1]):
            for j,ring2 in enumerate(mol.rings[i+1:]):
                biaryl_pivot = is_biaryl_ring(mol,ring1,ring2)
                if biaryl_pivot:
                    mol.biaryl_rings.append([i,i+j+1])
                    mol.biaryl_pivots.append(biaryl_pivot)
        if mol.option.verbose:
            print( 'Added regular biaryl-pivots: ')
            for a1,a2 in mol.biaryl_pivots:
                print( " (%s,%s)"%(mol.atms[a1].name,mol.atms[a2].name))

    ## 2. Pseudo ring-ring: non-ring connected to ring
    if len(mol.rings) > 0:
        biaryl_pivot_extra = []
        for i,ring1 in enumerate(mol.rings_aro):
            biaryl_pivot_extra += search_special_biaryl_ring(mol,ring1)
        if mol.option.verbose:
            print( 'Added extra biaryl-pivots: ')
            for a1,a2 in biaryl_pivot_extra:
                print( " (%s,%s)"%(mol.atms[a1].name,mol.atms[a2].name))
        mol.biaryl_pivots_extra = biaryl_pivot_extra
        mol.biaryl_pivots += biaryl_pivot_extra

    # 3. Amide bond; trick to allow rotation around amide bonds
    for i,bond in enumerate(mol.bonds):
        aclass1 = ACLASS_ID[mol.atms[bond.atm1].aclass]
        aclass2 = ACLASS_ID[mol.atms[bond.atm2].aclass]
        is_amide_connection = (aclass1 in ['Nad','Nad3']) and (aclass2 in ['Nad','Nad3'])
        if is_amide_connection: # append into special biaryl
            #A final check, if the initial conformation is planar around the amide bond
            #We need to treat it as conjugated bond and disallow the rotation.
            #GZ, Aug,20,2020
            if mol.option.opt.infer_conjugation and not is_planar(mol, bond.atm1, bond.atm2):
                mol.biaryl_pivots.append((bond.atm1,bond.atm2))

    if mol.option.opt.reassign_biaryl_aclass:
        reassign_biaryl_atypes(mol)
    else:
        assign_bond_conjugation(mol)

    if mol.option.verbose:
        print('Atom_Puckering: ', [mol.atms[atm].name for atm in mol.atms_puckering])
        print('Atom_Aro: ', [mol.atms[atm].name for atm in mol.atms_aro])
        print('Rings: ', [[mol.atms[atm].name for atm in ring.atms] for ring in mol.rings])
        if len(mol.biaryl_rings) > 0 :
            print('BiarylRings:', mol.biaryl_rings, ' BiarylAxes: ')
            for a1,a2 in mol.biaryl_pivots:
                print((mol.atms[a1].name,mol.atms[a2].name))
            print()

# Deprecated: For atom-type-based torsion assignment
def reassign_biaryl_atypes(mol):
    # reassign atype
    # Turned off as of Oct 2018 with new bond-based torsion assignment logic
    for a1,a2 in mol.biaryl_pivots:
        if mol.atms[a1].aclass in [ACLASS_ID.index('CR'),ACLASS_ID.index('CRp'),
                                   ACLASS_ID.index('CD1'),ACLASS_ID.index('CD2'),ACLASS_ID.index('CD'),ACLASS_ID.index('CDp')]: #carbon
            mol.atms[a1].aclass = ACLASS_ID.index('CRb')
        elif mol.atms[a1].atype == ATYPES.index('N'):
            mol.atms[a1].aclass = ACLASS_ID.index('NGb')

        if mol.atms[a2].aclass in [ACLASS_ID.index('CR'),ACLASS_ID.index('CRp'),
                                   ACLASS_ID.index('CD1'),ACLASS_ID.index('CD2'),ACLASS_ID.index('CDp'),ACLASS_ID.index('CDp')]: #carbon
            mol.atms[a2].aclass = ACLASS_ID.index('CRb')
        elif mol.atms[a2].atype == ATYPES.index('N'):
            mol.atms[a2].aclass = ACLASS_ID.index('NGb')

# For bond-type-based torsion assignment
# use bond-order + biaryl_pivot info
def assign_bond_conjugation(mol):
    for ibond,bond in enumerate(mol.bonds):
        # Consider not conjugated if any of two is sp3
        if mol.atms[bond.atm1].hyb == 3 or mol.atms[bond.atm2].hyb == 3: continue

        # Jan 2019: To avoid forcing planar forms for puckering rings
        # Also not conjugated if belongs to a puckering ring
        bond_at_puckering_ring = False
        for ring in mol.rings:
            if ring.type < 4 and ring.has((bond.atm1,bond.atm2)):
                bond_at_puckering_ring = True
                break
        if bond_at_puckering_ring: continue

        # special conjugation through ring(N=C)-N-H
        # THIS PART SHOULD BE REVISITED WITH "RING" BONDTYPE IN FUTURE
        is_ring_NCNH = False
        if (bond.atm1 in mol.atms_aro and bond.atm2 not in mol.atms_aro) or \
           (bond.atm2 in mol.atms_aro and bond.atm1 not in mol.atms_aro):
            # i: ring-member connected to other branch
            # j: non-ring-member connected to i
            if (bond.atm1 in mol.atms_aro and bond.atm2 not in mol.atms_aro):
                i,j = (bond.atm1,bond.atm2)
            else:
                i,j = (bond.atm2,bond.atm1)

            type_i = ATYPES[mol.atms[i].atype]
            type_j = ATYPES[mol.atms[j].atype]
            if mol.atms[j].has_H and (type_i,type_j) == ('C','N'):
                for k,dumm in mol.atms[i].bonds: #from C
                    if mol.in_same_ring(i,k) == 0 or k == j: continue
                    if (ATYPES[mol.atms[k].atype] == 'N') and (mol.atms[k].hyb != 3):
                        is_ring_NCNH = True
                        break

                #for l,dumm in mol.atms[j].bonds: #from N
                #    if l == i: continue
                #    if l in mol.atms_aro:
                #        is_ring_NCNH = False #override this rule if ringNC-NH-aro
                #        break

        if is_ring_NCNH:
            print("Torsion around %s-%s assigned as conjugated by [ring N=C]-N-H rule"%(mol.atms[bond.atm1].name,mol.atms[bond.atm2].name))
            mol.bonds[ibond].is_conjugated = True
            continue

        # also skip if is part of >= 7-membered ring (whether aromatic or not)
        # because conjugation start to break (more influenced by neighbors than conjugation)
        # THIS PART SHOULD BE REVISITED WITH "RING" BONDTYPE IN FUTURE
        minimum_ring_size = 100
        for ring in mol.rings:
            if ring.has((bond.atm1,bond.atm2)):
                if ring.natms < minimum_ring_size:
                    minimum_ring_size = ring.natms
        if minimum_ring_size < 100 and minimum_ring_size > 6:
            continue

        # simpler logic over-predicts conjugated rings
        #if mol.in_same_ring(bond.atm1,bond.atm2) == 1: continue

        #print( "assign conjugation: ", mol.atms[bond.atm1].name, mol.atms[bond.atm2].name,
        #       (bond.atm1,bond.atm2) in mol.biaryl_pivots or \
        #       (bond.atm2,bond.atm1) in mol.biaryl_pivots,
        #       ACLASS_ID[mol.atms[bond.atm1].aclass] in CONJUGATING_ACLASSES,
        #       ACLASS_ID[mol.atms[bond.atm2].aclass] in CONJUGATING_ACLASSES)

        #CONJUGATING_ACLASSES doesn't include Nad3, which in some structure
        #can form conjugate bond with aromatic rings.
        #Added Nad3 to CONJUGATING_ACLASSES. The conjugation will be determined
        #according to the planarity of the initial conformation.
        #GZ, Aug,20,2020
        if ACLASS_ID[mol.atms[bond.atm1].aclass] not in CONJUGATING_ACLASSES or \
           ACLASS_ID[mol.atms[bond.atm2].aclass] not in CONJUGATING_ACLASSES: continue

        if ((bond.atm1,bond.atm2) in mol.biaryl_pivots) or \
           ((bond.atm2,bond.atm1) in mol.biaryl_pivots):
           continue

        #Note: we need to be careful about assigning conjugated type to a single bond.
        # especially for (C=O)-NH-CR cases, in some cases NH is conjugated with C=O but
        # NH-CR is not conjugated. e.g. APAJIL from CSD.
        # Use the dihedral angle as a second check, if the angle is clearly not planar,
        # don't assign it as conjugated bond.
        # GZ, July,6,2020
        if bond.order == 1 and (not mol.atms[bond.atm1].is_H) \
            and (not mol.atms[bond.atm2].is_H):
            if mol.option.opt.infer_conjugation and not is_planar(mol, bond.atm1, bond.atm2):
                mol.bonds[ibond].is_conjugated = False #don't assign it conjugated
                continue

        #Note: any bond reaches at this point will be assigned as conjugated.
        #GZ, July,6,2020
        mol.bonds[ibond].is_conjugated = True

# Define "CHI"s
def define_rotable_torsions(mol):
    mol.chiatms = []
    mol.chitypes = []
    mol.chiextra = ""

    hapol_torsion_id = []
    hpol_torsion_type = [0 for k in mol.torsions]

    # 1. First get list of hydrogen torsions
    for i,torsion in enumerate(mol.torsions):
        aclasses = [ACLASS_ID[mol.atms[atm].aclass] for atm in torsion]
        # Apolar hydrogen torsions
        if (aclasses[0] in ACLASS_HAPOL) or (aclasses[3] in ACLASS_HAPOL):
            hapol_torsion_id.append(i)

        # PolarH torsions
        elif (aclasses[0] in ACLASS_HPOL) and (aclasses[3] not in ACLASS_HPOL):
            stem = mol.atms[torsion[1]]
            if stem.hyb == 2:
                if len(stem.bonds) < 3:
                    hpol_torsion_type[i] = 2
            elif stem.hyb == 3: #mol.torsion_orders[i] == 1:
                hpol_torsion_type[i] = 3

        elif (aclasses[3] in ACLASS_HPOL) and (aclasses[0] not in ACLASS_HPOL):
            stem = mol.atms[torsion[2]]
            if stem.hyb == 2:
                if len(stem.bonds) < 3:
                    hpol_torsion_type[i] = 2
            elif stem.hyb == 3: #mol.torsion_orders[i] == 1:
                hpol_torsion_type[i] = 3

    # 1-2. Get estimations of num_H_conf
    num_H_confs = 1
    covered = []
    for i,atms in enumerate(mol.torsions):
        if (atms[1],atms[2]) in covered or (atms[2],atms[1]) in covered: continue
        if hpol_torsion_type[i] == 2: num_H_confs *= 6
        elif hpol_torsion_type[i] == 3: num_H_confs *= 9
        covered.append((atms[1],atms[2]))

    #print 'Total num_H_conf: ', num_H_confs
    if num_H_confs <= mol.max_confs: mol.chiextra = "1 20"
    else: mol.chiextra = "0"

    # 2. CHI assignment; read-in from mol.option
    # mol.torsions SHOULD have been ordered such that
    # heavyatom-only comes first followed by H-containing ones
    ring_cuts = []
    for ring in mol.rings: ring_cuts.append(ring.cut_bond)

    covered = []
    for i,atms in enumerate(mol.torsions):
        atm0 = mol.atms[atms[0]]
        atm1 = mol.atms[atms[1]]
        atm2 = mol.atms[atms[2]]
        atm3 = mol.atms[atms[3]]
        #GZ, make torsion around CSQ @ CSQ non-rotatable, Sep 3, 2020
        if mol.in_same_ring(atms[1],atms[2]) and ACLASS_ID[atm1.aclass] == "CSQ" and ACLASS_ID[atm2.aclass] == "CSQ":
            continue
        border = mol.bond_order(atms[1],atms[2])

        if (atms[1],atms[2]) in covered or (atms[2],atms[1]) in covered: continue #avoid
        if ((atm0.root != atms[1]) and (atm1.root != atms[0])) or ((atm2.root != atms[3]) and (atm3.root != atms[2])): continue

        if atm1.root == atms[2]:
            atms_ordered = [atms[3],atms[2],atms[1],atms[0]] #last is tipatm
        elif atm2.root == atms[1]:
            atms_ordered = [atms[0],atms[1],atms[2],atms[3]] #last is tipatm
        else:
            continue

        # avoid if any of bonds defined as cut_bond
        if (atms[0],atms[1]) in ring_cuts or (atms[1],atms[0]) in ring_cuts or \
           (atms[1],atms[2]) in ring_cuts or (atms[2],atms[1]) in ring_cuts or \
           (atms[2],atms[3]) in ring_cuts or (atms[3],atms[2]) in ring_cuts:
            continue

        # why did I put this logic here...?
        # check if the torsion is non-ATorder but connected to ATorder
        FT_connected = False
        for j in range(len(atms)-1):
            a1,a2 = atms[j],atms[j+1]
            if a1 in mol.ATorder and a2 in mol.ATorder: continue
            if (mol.atms[a1].root not in [a2,a1]) and (mol.atms[a2].root not in [a1,a2]):
                FT_connected = True
                break
        if FT_connected: continue

        covered.append((atms[1],atms[2]))

        is_biaryl_pivot = ((atms[1],atms[2]) in mol.biaryl_pivots) or \
                          ((atms[2],atms[1]) in mol.biaryl_pivots)

        # Bug fix with biaryl not defined as CHI: Oct17 2018
        if ((atms[1] in mol.atms_aro) and (atms[2] in mol.atms_aro)):
            if border>1: continue
            same_ring_order = mol.in_same_ring(atms[1],atms[2])
            if same_ring_order == 4:
                continue # also skip if belongs to the same "aromatic" ring
            elif same_ring_order == 0: #not in same ring
                if not mol.option.opt.report_ringring_chi and is_biaryl_pivot:
                    continue

        # below are optional
        # Apolar hydrogen chis
        if (not mol.option.opt.report_Hapol_chi) and (i in hapol_torsion_id): continue

        if (mol.bond_order(atms[1],atms[2]) > 1):
            atype1 = ACLASS_ID[atm1.aclass]
            atype2 = ACLASS_ID[atm2.aclass]
            is_amide_bond = ((atype1 == 'Nad') and (atype2 in ['CDp','CRb'])) or \
                            ((atype2 == 'Nad') and (atype1 in ['CDp','CRb']))

            #is_aliphatic_bond = ((atype1 == 'CD1') and (atype2 == 'CD1') and \
            #                     (atms[1] not in mol.atms_ring) and (atms[2] not in mol.atms_ring))

            #print atype1, atype2, is_amide_bond
            #print( mol.atms[atms[1]].name,mol.atms[atms[2]].name,is_aliphatic_bond)
            #if is_aliphatic_bond:
            #    pass

            ## always skip a bond in a same aromatic ring
            if mol.in_same_ring(atms[1],atms[2]):
                continue

            if (mol.option.opt.report_amide_chi):
                if (not is_amide_bond): continue
            elif (not mol.option.opt.report_nbonded_chi and not is_biaryl_pivot):
                continue

        # skip conjugated polarH chi
        elif mol.bond_conjugated(atms[1],atms[2]):
            nH1 = 0
            for j,order in atm1.bonds:
                if mol.atms[j].is_H: nH1+=1
            nH2 = 0
            for j,order in atm2.bonds:
                if mol.atms[j].is_H: nH2+=1

            if nH1 == len(atm1.bonds)-1 or nH2 == len(atm2.bonds)-1:
                continue

        if (not mol.option.opt.report_puckering_chi) and \
           (atms[1] in mol.atms_puckering) and (atms[2] in mol.atms_puckering):
            continue

        mol.chiatms.append(atms_ordered)

        if hpol_torsion_type[i] == 2: #sp2
            mol.chitypes.append('sp2')
        elif hpol_torsion_type[i] == 3: #sp3
            mol.chitypes.append('sp3')
        elif i in hapol_torsion_id: #sp3
            mol.chitypes.append('sp3H')
        #elif (atms[1] in mol.atms_puckering) and (atms[2] in mol.atms_puckering):
        #    mol.chitypes.append('pucker') #turn this functionality for now... add separate flag for this behavior
        else:
            mol.chitypes.append('')


import re
import math
import numpy
try:
    from galpy.util import bovy_plot
except ImportError:
    import bovy_plot #BOVY: COPY IN LOCAL VERSION
_ZSOLAR= 0.019
_LOGTESUN= numpy.log10(5777)
_LOGGSUN= numpy.log10(27400.)
class Isochrone:
    """Template for any Isochrone type class"""
    def __init__(self):
        """
        NAME:
           __init__
        PURPOSE:
           initialize
        INPUT:
        OUTPUT:
        HISTORY:
           2011-04-27 - Written - Bovy (NYU)
        """
        raise NotImplementedError("'__init__' not implemented for this isochrone")

    def __call__(self,logage,Z=None,feh=None,afe=None,maxm=None,
                 stage=None):
        """
        NAME:
           __call__
        PURPOSE:
           get a single isochrone from the library
        INPUT:
           logage - log_10 age
           Z= or feh= metallicity (use Z_\odot=0.019)
           afe= None (not supported for Padova)
           maxm= maximum mass to consider (m_ini)
           stage= if set, only return this evolutionary stage
                  (if this exists for this isochrone libary)
        OUTPUT:
           isochrone
        HISTORY:
           Written - Bovy (NYU)
        """
        raise NotImplementedError("'__call__' method for this isochrone is not implemented")

    def Zs(self):
        """
        NAME:
           Zs
        PURPOSE:
           Return the loaded metallicities
        INPUT:
        OUTPUT:
        HISTORY:
           2011-04-27 - Written - Bovy (NYU)
        """
        return self._ZS

    def logages(self):
        """
        NAME:
           logages
        PURPOSE:
           Return the loaded log_10 ages
        INPUT:
        OUTPUT:
        HISTORY:
           2011-04-27 - Written - Bovy (NYU)
        """
        return self._logages

    def filters(self):
        """
        NAME:
           filters
        PURPOSE:
           Return the supported filters
        INPUT:
        OUTPUT:
        HISTORY:
           2011-04-27 - Written - Bovy (NYU)
        """
        return self._filters
###################################PLOTTING####################################
    def plot(self,logage,*args,**kwargs):
        """
        NAME:
           plot
        PURPOSE:
           plot an individual isochrone, or a set of isochrones
        INPUT:
           logage - log_10 age or list thereof
           Z= metallicity or list
           feh= metallicity or list (use Z_\odot=0.019)
           afe= not supported for Padova
           d1= x dimension (for color write 'J-Ks')
           d2= y dimension
           maxm= maximum mass to plot
           stage= if set, only return this evolutionary stage
                  (if this exists for this isochrone libary)
           ignore_gaps= if True, ignore non-existant isochrones
           +bovy_plot.bovy_plot keywords
        OUTPUT:
           plot to output device
        HISTORY:
           2011-04-27 - Written - Bovy (NYU)
        """
        if not isinstance(logage,(list,numpy.ndarray)) \
                and not (('Z' in kwargs \
                              and isinstance(kwargs['Z'],(list,numpy.ndarray)))\
                             or ('feh' in kwargs \
                                     and isinstance(kwargs['feh'],numpy.ndarray))):
            return self._plot_single(logage,*args,**kwargs)
        #Do we have Z or FeH?
        if not 'Z' in kwargs and 'feh' in kwargs: usefeh= True
        else: usefeh= False
        if not isinstance(logage,(list,numpy.ndarray)) and usefeh:
            logage= numpy.array([logage for ii in range(len(kwargs['feh']))])
        elif not isinstance(logage,(list,numpy.ndarray)):
            logage= numpy.array([logage for ii in range(len(kwargs['Z']))])
        #Handle Z etc.
        if 'feh' in kwargs:
            if isinstance(kwargs['feh'],(list,numpy.ndarray)):
                fehs= kwargs['feh']
            else:
                fehs= numpy.array([kwargs['feh'] for ii in range(len(logage))])
        else:
            fehs= [None for ii in range(len(logage))]
        if 'Z' in kwargs:
            if isinstance(kwargs['Z'],(list,numpy.ndarray)):
                ZS= kwargs['Z']
            else:
                ZS= numpy.array([kwargs['Z'] for ii in range(len(logage))])
        else: ZS= [None for ii in range(len(logage))]
        if usefeh and not len(logage) == len(fehs):
            raise IOError("When both logage and feh are given as arrays they need to have the same length")
        elif not usefeh and not len(logage) == len(ZS):
            raise IOError("When both logage and Z are given as arrays they need to have the same length")
        #Plot first
        if usefeh: kwargs['feh']= fehs[0]
        else: kwargs['Z']= ZS[0]
        out= self._plot_single(logage[0],*args,**kwargs)
        kwargs['overplot']= True
        for ii in range(1,len(logage)):
            kwargs['Z']= ZS[ii]
            kwargs['feh']= fehs[ii]
            self._plot_single(logage[ii],*args,**kwargs)
        return out

    def _plot_single(self,logage,*args,**kwargs):
        #kwargs
        Z= kwargs.pop('Z',None)
        feh= kwargs.pop('feh',None)
        afe= kwargs.pop('afe',None)
        maxm= kwargs.pop('maxm',None)
        stage= kwargs.pop('stage',None)
        d1= kwargs.pop('d1',self._filters[0]+'-'+self._filters[1])
        d2= kwargs.pop('d2',self._filters[0])
        ignore_gaps= kwargs.pop('ignore_gaps',False)
        #get isochrone
        try:
            iso= self(logage,Z=Z,feh=feh,afe=afe,maxm=maxm,stage=stage)
        except IOError:
            if ignore_gaps: return None
            else:
                raise IOError("No isochrone found for this logage/metallicity combination\nUse ignore_gaps=True to ignore non-existant isochrones")
        #get dimensions
        colorx= re.split(r'-',d1)
        if len(colorx) == 2: #d1 is a color
            x= iso[colorx[0]]-iso[colorx[1]]
        else:
            x= iso[d1]
        colory= re.split(r'-',d2)
        if len(colory) == 2: #d2 is a color
            y= iso[colory[0]]-iso[colory[1]]
        else:
            y= iso[d2]
        #Put in default labels
        if not kwargs.get('overplot',False):
            kwargs['xlabel']= kwargs.get('xlabel',r'$'+d1+'$')
            kwargs['ylabel']= \
                kwargs.get('ylabel',
                    r'$M_{'+d2+'}$' if d2 in self._filters else r'$'+d2+'$')
            if not 'yrange' in kwargs and d2 in self._filters:
                kwargs['yrange']= [numpy.amax(y)+0.3,numpy.amin(y)-0.3]
        #plot
        return bovy_plot.bovy_plot(x,y,*args,**kwargs)

def Z2FEH(z,zsolar=None,parsec=False):
    """Convert Z to FeH assuming zsolar"""
    if parsec:
        if zsolar is None: zsolar= 0.0152
        return numpy.log10(z/(1.-0.2485-2.78*z))-math.log10(zsolar/(1.-0.2485-2.78*zsolar))
    else:
        if zsolar is None: zsolar= _ZSOLAR
        return numpy.log10(z)-math.log10(zsolar)

def FEH2Z(feh,zsolar=None,parsec=False):
    """Convert FeH to Z assuming zsolar"""
    if parsec:
        if zsolar is None: zsolar= 0.0152
        zx= 10.**(feh+math.log10(zsolar/(1.-0.2485-2.78*zsolar)))
        return (zx-0.2485*zx)/(2.78*zx+1.)
    else:
        if zsolar is None: zsolar= _ZSOLAR
        return 10.**(feh+math.log10(zsolar))

def logg(logL,logTe,mass):
    """
    NAME:
       logg
    PURPOSE:
       calculate log g from luminosity, teff, and mass
    INPUT:
       logL - log Luminosity (solar units)
       logTe- log effective temperature (log K)
       mass - mass (solar units)
    OUTPUT:
       log g (log cm/s^2)
    HISTORY:
       2012-08-16 - Written - Bovy (IAS)
    """
    logR= -2.*(logTe-_LOGTESUN)+0.5*logL
    return numpy.log10(mass)-2.*logR+_LOGGSUN

def dict2recarray(dict):
    nEntries= len(dict.keys())
    nOut= len(dict[dict.keys()[0]])
    out= numpy.zeros(nOut,dtype={'names':dict.keys(),
                                 'formats':[numpy.float64 for ii in range(nEntries)]})
    for ii in range(nEntries):
        out[dict.keys()[ii]]= dict[dict.keys()[ii]]
    return out.view(numpy.recarray)