import abc
from collections import defaultdict
from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union, overload

import numpy
from openff.utilities import requires_package
from pydantic import BaseModel, Field, constr
from typing_extensions import Literal

from openff.recharge.aromaticity import AromaticityModel, AromaticityModels
from openff.recharge.charges.exceptions import UnableToAssignChargeError
from openff.recharge.utilities.openeye import import_oechem

if TYPE_CHECKING:

    try:
        import torch
    except ImportError:
        torch = None

    from openeye.oechem import OEMol
    from openff.toolkit.topology import Molecule
    from openff.toolkit.typing.engines.smirnoff import VirtualSiteHandler

ExclusionPolicy = Literal["none", "parents"]

VirtualSiteKey = Tuple[str, str, str]
VirtualSiteChargeKey = Tuple[str, str, str, int]
VirtualSiteGeometryKey = Tuple[
    str, str, str, Literal["distance", "in_plane_angle", "out_of_plane_angle"]
]

_DEGREES_TO_RADIANS = numpy.pi / 180.0

class _VirtualSiteParameter(BaseModel, abc.ABC):
    """The base class for virtual site parameters."""

    type: Literal["base-virtual-site"]

    smirks: constr(min_length=1) = Field(
        ...,
        description="A SMIRKS pattern that encodes the chemical environment that "
        "this parameter should be applied to.",
    )
    name: Optional[str] = Field(
        None, description="An optional name associated with this virtual site."
    )

    distance: float = Field(
        ...,
        description="The distance to place the virtual site along its associated basis.",
    )
    charge_increments: Tuple[float, ...] = Field(
        ...,
        description="The amount of charge [e] to be transferred from the virtual site "
        "to each tagged atom that forms the basis for the virtual site.",
    )

    sigma: float = Field(
        ..., description="The LJ sigma parameter [A] associated with the virtual site."
    )
    epsilon: float = Field(
        ...,
        description="The LJ espilon [kJ / mol] parameter associated with the virtual "
        "site.",
    )

    match: Literal["once", "all-permutations"] = Field(..., description="...")

    @classmethod
    @abc.abstractmethod
    def local_frame_weights(cls) -> numpy.ndarray:
        """Returns a matrix of the weights to apply to a matrix of the coordinates of
        the virtual sites' parent atoms to yield the origin, x and y vectors of the
        virtual sites local frame with shape=(3, n_parent_atoms)."""
        raise NotImplementedError()

    @property
    @abc.abstractmethod
    def local_frame_coordinates(self) -> numpy.ndarray:
        """Returns a 1 X 3 array of the spherical coordinates (``[d, theta, phi]``) of
        this virtual site with respect to its local frame.
        """
        raise NotImplementedError()

class BondChargeSiteParameter(_VirtualSiteParameter):

    type: Literal["BondCharge"] = "BondCharge"

    @classmethod
    def local_frame_weights(cls) -> numpy.ndarray:
        return numpy.array([[1.0, 0.0], [-1.0, 1.0], [-1.0, 1.0]])

    @property
    def local_frame_coordinates(self) -> numpy.ndarray:
        # distance, theta, phi
        return numpy.array([[self.distance, 180.0, 0.0]])

class MonovalentLonePairParameter(_VirtualSiteParameter):

    type: Literal["MonovalentLonePair"] = "MonovalentLonePair"

    in_plane_angle: float = Field(
        ...,
        description="The angle [deg] to move the virtual site in the plane defined "
        "by the tagged atoms by.",
    )
    out_of_plane_angle: float = Field(
        ...,
        description="The angle [deg] to move the virtual site out of the plane "
        "defined by the tagged atoms by.",
    )

    @classmethod
    def local_frame_weights(cls) -> numpy.ndarray:
        return numpy.array([[1.0, 0.0, 0.0], [-1.0, 1.0, 0.0], [-1.0, 0.0, 1.0]])

    @property
    def local_frame_coordinates(self) -> numpy.ndarray:
        # distance, theta, phi
        return numpy.array(
            [[self.distance, self.in_plane_angle, self.out_of_plane_angle]]
        )

class DivalentLonePairParameter(_VirtualSiteParameter):

    type: Literal["DivalentLonePair"] = "DivalentLonePair"

    out_of_plane_angle: float = Field(
        ...,
        description="The angle [deg] to move the virtual site out of the plane "
        "defined by the tagged atoms by.",
    )

    @classmethod
    def local_frame_weights(cls) -> numpy.ndarray:
        return numpy.array([[0.0, 1.0, 0.0], [0.5, -1.0, 0.5], [1.0, -1.0, 0.0]])

    @property
    def local_frame_coordinates(self) -> numpy.ndarray:
        # distance, theta, phi
        return numpy.array([[self.distance, 180.0, self.out_of_plane_angle]])

class TrivalentLonePairParameter(_VirtualSiteParameter):

    type: Literal["TrivalentLonePair"] = "TrivalentLonePair"

    @classmethod
    def local_frame_weights(cls) -> numpy.ndarray:
        return numpy.array(
            [
                [0.0, 1.0, 0.0, 0.0],
                [1.0 / 3.0, -1.0, 1.0 / 3.0, 1.0 / 3.0],
                [1.0, -1.0, 0.0, 0.0],
            ]
        )

    @property
    def local_frame_coordinates(self) -> numpy.ndarray:
        # distance, theta, phi
        return numpy.array([[self.distance, 180.0, 0.0]])

VirtualSiteParameterType = Union[
    BondChargeSiteParameter,
    MonovalentLonePairParameter,
    DivalentLonePairParameter,
    TrivalentLonePairParameter,
]

class VirtualSiteCollection(BaseModel):
    """A collection of virtual site parameters that are based off of the SMIRNOFF
    specification."""

    parameters: List[VirtualSiteParameterType] = Field(
        ...,
        description="The virtual site parameters to apply.",
    )
    aromaticity_model: AromaticityModels = Field(
        AromaticityModels.MDL,
        description="The model to use when assigning aromaticity.",
    )

    exclusion_policy: ExclusionPolicy = Field("parents", description="...")

    @requires_package("openff.toolkit")
    @requires_package("simtk")
    def to_smirnoff(self) -> "VirtualSiteHandler":
        """Converts this collection of virtual site parameters to a SMIRNOFF virtual
        site parameter handler.

        Returns
        -------
            The constructed parameter handler.
        """

        from openff.toolkit.typing.engines.smirnoff import VirtualSiteHandler
        from simtk import unit

        # noinspection PyTypeChecker
        parameter_handler = VirtualSiteHandler(
            version="0.3", exclusion_policy=self.exclusion_policy
        )

        for parameter in reversed(self.parameters):

            parameter_kwargs = dict(
                smirks=parameter.smirks,
                type=parameter.type,
                name=parameter.name,
                distance=parameter.distance * unit.angstrom,
                charge_increment=[
                    charge * unit.elementary_charge
                    for charge in parameter.charge_increments
                ],
                sigma=parameter.sigma * unit.angstrom,
                epsilon=parameter.epsilon * unit.kilojoules_per_mole,
                match=parameter.match.replace("-", "_").lower(),
            )

            if parameter.type == "MonovalentLonePair":

                parameter_kwargs["outOfPlaneAngle"] = (
                    parameter.out_of_plane_angle * unit.degrees
                )
                parameter_kwargs["inPlaneAngle"] = (
                    parameter.in_plane_angle * unit.degrees
                )

            elif parameter.type == "DivalentLonePair":

                parameter_kwargs["outOfPlaneAngle"] = (
                    parameter.out_of_plane_angle * unit.degrees
                )

            parameter_handler.add_parameter(parameter_kwargs=parameter_kwargs)

        return parameter_handler

    @classmethod
    @requires_package("simtk")
    def from_smirnoff(
        cls,
        parameter_handler: "VirtualSiteHandler",
        aromaticity_model=AromaticityModels.MDL,
    ) -> "VirtualSiteCollection":
        """Attempts to convert a SMIRNOFF virtual site parameter handler to a virtual
        site parameter collection.

        Parameters
        ----------
        parameter_handler
            The parameter handler to convert.
        aromaticity_model
            The model which describes how aromaticity should be assigned
            when applying the virtual site correction parameters.

        Returns
        -------
            The converted virtual site collection.
        """

        from simtk import unit

        parameters = []

        for smirnoff_parameter in reversed(parameter_handler.parameters):

            base_kwargs = dict(
                smirks=smirnoff_parameter.smirks,
                name=smirnoff_parameter.name,
                distance=smirnoff_parameter.distance.value_in_unit(unit.angstrom),
                charge_increments=tuple(
                    charge.value_in_unit(unit.elementary_charge)
                    for charge in smirnoff_parameter.charge_increment
                ),
                sigma=smirnoff_parameter.sigma.value_in_unit(unit.angstrom),
                epsilon=smirnoff_parameter.epsilon.value_in_unit(
                    unit.kilojoules_per_mole
                ),
                match=smirnoff_parameter.match.replace("_", "-").lower(),
            )

            if smirnoff_parameter.type == "BondCharge":
                parameter = BondChargeSiteParameter(**base_kwargs)

            elif smirnoff_parameter.type == "MonovalentLonePair":

                parameter = MonovalentLonePairParameter(
                    **base_kwargs,
                    out_of_plane_angle=smirnoff_parameter.outOfPlaneAngle.value_in_unit(
                        unit.degrees
                    ),
                    in_plane_angle=smirnoff_parameter.inPlaneAngle.value_in_unit(
                        unit.degrees
                    ),
                )

            elif smirnoff_parameter.type == "DivalentLonePair":

                parameter = DivalentLonePairParameter(
                    **base_kwargs,
                    out_of_plane_angle=smirnoff_parameter.outOfPlaneAngle.value_in_unit(
                        unit.degrees
                    ),
                )

            elif smirnoff_parameter.type == "TrivalentLonePair":
                parameter = TrivalentLonePairParameter(**base_kwargs)

            else:
                raise NotImplementedError()

            parameters.append(parameter)

        return VirtualSiteCollection(
            parameters=parameters,
            aromaticity_model=aromaticity_model,
            exclusion_policy=parameter_handler.exclusion_policy.lower(),
        )

    def vectorize_coordinates(
        self, parameter_keys: List[VirtualSiteGeometryKey]
    ) -> numpy.ndarray:
        """Returns a flat vector of the local frame coordinate values associated with a
        specified set of 'keys'.

        Parameters
        ----------
        parameter_keys
            A list of parameter 'keys' of the form ``(smirks, type, name, attr)`` that
            specify which local frame coordinate to include in the returned vector.

            The allowed attributes are ``distance``, ``in_plane_angle``,
            ``out_of_plane_angle``

        Returns
        -------
            A vector of local frame coordinate with shape=(n_keys, 1)
        """

        parameters_by_key = {
            (parameter.smirks, parameter.type, parameter.name): parameter
            for parameter in self.parameters
        }

        parameter_values = numpy.array(
            [
                [getattr(parameters_by_key[tuple(parameter_key)], attribute)]
                for *parameter_key, attribute in parameter_keys
            ]
        )

        return parameter_values

    def vectorize_charge_increments(
        self, parameter_keys: List[VirtualSiteChargeKey]
    ) -> numpy.ndarray:
        """Returns a flat vector of the charge increment values associated with a
        specified set of 'keys'.

        Parameters
        ----------
        parameter_keys
            A list of parameter 'keys' of the form ``(smirks, type, name, idx)`` that
            specify which charge increments to include in the returned vector, where
            `idx` is an integer index into a parameters' ``charge_increments`` tuple.

        Returns
        -------
            A vector of charge increments with shape=(n_keys, 1)
        """

        parameters_by_key = {
            (parameter.smirks, parameter.type, parameter.name): parameter
            for parameter in self.parameters
        }

        return numpy.array(
            [
                [
                    parameters_by_key[tuple(parameter_key)].charge_increments[
                        charge_index
                    ]
                ]
                for *parameter_key, charge_index in parameter_keys
            ]
        )

class VirtualSiteGenerator:
    @classmethod
    def _apply_virtual_sites(
        cls, oe_molecule: "OEMol", vsite_collection: VirtualSiteCollection
    ) -> Tuple["Molecule", Dict[Tuple[int, ...], List[VirtualSiteKey]]]:
        """Applies a virtual site collection to a molecule.

        Parameters
        ----------
        oe_molecule
            The molecule to build the virtual sites for.
        vsite_collection
            The v-site collection to use to create the virtual sites.

        Returns
        -------
            An OpenFF molecule with virtual sites as well as a dictionary that maps each
            virtual site back to the parameter that yielded it.
        """

        from openff.toolkit.topology import Molecule

        parameter_handler = vsite_collection.to_smirnoff()

        off_topology = parameter_handler.create_openff_virtual_sites(
            Molecule.from_openeye(oe_molecule).to_topology()
        )
        off_molecule = next(off_topology.reference_molecules)

        term_topology = parameter_handler._term_map_topology(off_topology)

        assigned_vsite_keys = defaultdict(set)

        for (_, atom_indices), vsite_keys in term_topology.items():

            for vsite_key in vsite_keys:

                (_, (smirks, (vsite_type, (vsite_name, _)))) = vsite_key
                assigned_vsite_keys[atom_indices].add((smirks, vsite_type, vsite_name))

        return off_molecule, {
            atom_indices: [*keys] for atom_indices, keys in assigned_vsite_keys.items()
        }

    @classmethod
    def _build_charge_increment_array(
        cls, vsite_collection: VirtualSiteCollection
    ) -> Tuple[numpy.ndarray, List[VirtualSiteChargeKey]]:
        """Returns a flat vector of the charge increments contained within a virtual site
        collection as well as a list of keys that map each value back to its original
        parameter.

        Parameters
        ----------
        vsite_collection
            The collection containing the v-site parameters.

        Returns
        -------
            ...
        """

        charge_values = []
        charge_keys = []

        for parameter in vsite_collection.parameters:

            for i, charge_increment in enumerate(parameter.charge_increments):

                charge_values.append(charge_increment)

                charge_keys.append(
                    (parameter.smirks, parameter.type, parameter.name, i)
                )

        return numpy.array(charge_values), charge_keys

    @classmethod
    def _validate_charge_assignment_matrix(
        cls,
        assignment_matrix: numpy.ndarray,
    ):
        """Validates the charge increment assignment matrix.

        Parameters
        ----------
        assignment_matrix
            The assignment matrix to validate with
            shape=(n_atoms + n_vsites, n_charge_increments)
        """

        non_zero_assignments = assignment_matrix.sum(axis=0).astype(bool)

        if non_zero_assignments.any():

            raise UnableToAssignChargeError(
                "An internal error occurred. The v-site charge increments alter the "
                "total charge of the molecule"
            )

    @classmethod
    def build_charge_assignment_matrix(
        cls,
        oe_molecule: "OEMol",
        vsite_collection: VirtualSiteCollection,
    ) -> numpy.ndarray:
        """Generates a matrix that specifies which v-site charge increments have been
        applied to which atoms in the molecule.

        The matrix takes the form ...

        Parameters
        ----------
        oe_molecule
            The molecule to assign the v-site charge increments to.
        vsite_collection
            The v-site parameters that may be assigned.

        Returns
        -------
            The assignment matrix with shape=(n_atoms + n_vsites, n_charge_increments)
            where ...
        """

        oechem = import_oechem()

        # Make a copy of the molecule to assign the aromatic flags to.
        oe_molecule = oechem.OEMol(oe_molecule)
        # Assign aromaticity flags to ensure correct smirks matches.
        AromaticityModel.assign(oe_molecule, vsite_collection.aromaticity_model)

        off_molecule, assigned_vsite_keys = cls._apply_virtual_sites(
            oe_molecule, vsite_collection
        )

        _, all_vsite_keys = cls._build_charge_increment_array(vsite_collection)

        n_particles = off_molecule.n_particles
        n_corrections = len(all_vsite_keys)

        assignment_matrix = numpy.zeros((n_particles, n_corrections))

        for vsite_particle in [
            vsite_particle
            for vsite in off_molecule.virtual_sites
            for vsite_particle in vsite.particles
        ]:

            vsite_index = vsite_particle.molecule_particle_index
            vsite_parameter_keys = assigned_vsite_keys[vsite_particle.orientation]

            for vsite_parameter_key in vsite_parameter_keys:

                for i, atom_index in enumerate(vsite_particle.orientation):

                    # noinspection PyTypeChecker
                    vsite_parameter_index = all_vsite_keys.index(
                        (*vsite_parameter_key, i)
                    )

                    assignment_matrix[atom_index, vsite_parameter_index] += 1
                    assignment_matrix[vsite_index, vsite_parameter_index] -= 1

        cls._validate_charge_assignment_matrix(assignment_matrix)
        return assignment_matrix

    @classmethod
    def apply_charge_assignment_matrix(
        cls,
        assignment_matrix: numpy.ndarray,
        vsite_collection: VirtualSiteCollection,
    ) -> numpy.ndarray:
        """Applies an assignment matrix to a list of virtual site parameters to yield the
        final charges increments due to the virtual sites for a molecule.

        Parameters
        ----------
        assignment_matrix
            The virtual site charge increment assignment matrix constructed using
            ``build_charge_assignment_matrix`` that describes how the virtual site
            charge increments should be applied. This should have
            shape=(n_atoms + n_vsites, n_charge_increments)
        vsite_collection
            The virtual site parameters that may be assigned.

        Returns
        -------
            The charge increments with shape=(n_atoms + n_vsites, 1).
        """

        correction_values, _ = cls._build_charge_increment_array(vsite_collection)
        charge_corrections = assignment_matrix @ correction_values

        if not numpy.isclose(charge_corrections.sum(), 0.0):

            raise UnableToAssignChargeError(
                "An internal error occurred. The bond charge corrections were applied "
                "in such a way so that the total charge of the molecule will be "
                "altered."
            )

        return charge_corrections.reshape(-1, 1)

    @classmethod
    def generate_charge_increments(
        cls,
        oe_molecule: "OEMol",
        vsite_collection: VirtualSiteCollection,
    ) -> numpy.ndarray:
        """Generate a set of charge increments due to virtual sites for a molecule.

        Parameters
        ----------
        oe_molecule
            The molecule to generate the charge increments for.
        vsite_collection
            The virtual site parameters that may be assigned.

        Returns
        -------
            The charge increments with shape=(n_atoms + n_vsites, 1) that should be
            applied to the molecule.
        """

        assignment_matrix = cls.build_charge_assignment_matrix(
            oe_molecule, vsite_collection
        )

        generated_corrections = cls.apply_charge_assignment_matrix(
            assignment_matrix, vsite_collection
        )

        return generated_corrections

    @classmethod
    def build_local_coordinate_frames(
        cls,
        conformer: numpy.ndarray,
        assigned_parameters: Dict[Tuple[int, ...], List[VirtualSiteParameterType]],
    ) -> numpy.ndarray:
        """Builds an orthonormal coordinate frame for each virtual particle
        based on the type of virtual site and the coordinates of the parent atoms.

        Notes
        -----
        * See `the OpenMM documentation for further information
          <http://docs.openmm.org/7.0.0/userguide/theory.html#virtual-sites>`_.

        Parameters
        ----------
        conformer
            The conformer of the molecule that the virtual sites are being added to
            with shape=(n_atoms, 3) and units of [A].
        assigned_parameters
            A dictionary of the form ``assigned_parameters[atom_indices] = parameters``
            where ``atom_indices`` is a tuple of indices corresponding to the atoms
            that the virtual site is connected to, and ``parameters`` is a list of the
            parameters that describe the virtual sites.

        Returns
        -------
            An array storing the local frames of all virtual sites with
            shape=(4, n_vsites, 3) whereby ``local_frames[0]`` is an array of the
            origins of each frame, ``local_frames[1]`` the x-directions,
            ``local_frames[2]`` the y-directions, and ``local_frames[2]`` the
            z-directions.
        """

        stacked_frames = [[], [], [], []]

        for parent_indices, vsite_parameter in (
            (parent_indices, vsite_parameter)
            for parent_indices, vsite_parameters in assigned_parameters.items()
            for vsite_parameter in vsite_parameters
        ):
            parent_coordinates = conformer[parent_indices, :]

            weighted_coordinates = (
                vsite_parameter.local_frame_weights() @ parent_coordinates
            )

            origin = weighted_coordinates[0, :]

            xy_plane = weighted_coordinates[1:, :]
            xy_plane_norm = xy_plane / numpy.sqrt(
                (xy_plane * xy_plane).sum(-1)
            ).reshape(-1, 1)

            x_hat = xy_plane_norm[0, :]
            z_hat = numpy.cross(x_hat, xy_plane[1, :])
            y_hat = numpy.cross(z_hat, x_hat)

            stacked_frames[0].append(origin.reshape(1, -1))
            stacked_frames[1].append(x_hat.reshape(1, -1))
            stacked_frames[2].append(y_hat.reshape(1, -1))
            stacked_frames[3].append(z_hat.reshape(1, -1))

        local_frames = numpy.stack([numpy.vstack(frames) for frames in stacked_frames])
        return local_frames

    @classmethod
    @overload
    def convert_local_coordinates(
        cls,
        local_frame_coordinates: numpy.ndarray,
        local_coordinate_frames: numpy.ndarray,
        backend: Literal["numpy"],
    ) -> numpy.ndarray:
        ...

    @classmethod
    @overload
    def convert_local_coordinates(
        cls,
        local_frame_coordinates: "torch.Tensor",
        local_coordinate_frames: "torch.Tensor",
        backend: Literal["torch"],
    ) -> "torch.Tensor":
        ...

    @classmethod
    def convert_local_coordinates(
        cls,
        local_frame_coordinates,
        local_coordinate_frames,
        backend: Literal["numpy", "torch"] = "numpy",
    ) -> numpy.ndarray:
        """Converts a set of local virtual site coordinates defined in a spherical
        coordinate system into a full set of cartesian coordinates.

        Parameters
        ----------
        local_frame_coordinates
            An array containing the local coordinates with shape=(n_vsites, 3) and with
            columns of distance [A], 'in plane angle' [deg] and 'out of plane'
            angle [deg].
        local_coordinate_frames
            The orthonormal basis associated with each of the virtual sites with
            shape=(4, n_vsites, 3). See the ``build_local_coordinate_frames`` function
            for more details.
        backend
            The framework to use when performing mathematical operations.

        Returns
        -------
            An array of the cartesian coordinates of the virtual sites with
            shape=(n_vsites, 3) and units of [A].
        """

        if backend == "numpy":
            np = numpy
        elif backend == "torch":
            import torch

            np = torch
        else:
            raise NotImplementedError()

        d = local_frame_coordinates[:, 0].reshape(-1, 1)

        theta = (local_frame_coordinates[:, 1] * _DEGREES_TO_RADIANS).reshape(-1, 1)
        phi = (local_frame_coordinates[:, 2] * _DEGREES_TO_RADIANS).reshape(-1, 1)

        cos_theta = np.cos(theta)
        sin_theta = np.sin(theta)

        cos_phi = np.cos(phi)
        sin_phi = np.sin(phi)

        # Here we use cos(phi) in place of sin(phi) and sin(phi) in place of cos(phi)
        # this is because we want phi=0 to represent a 0 degree angle from the x-y plane
        # rather than 0 degrees from the z-axis.
        vsite_positions = (
            local_coordinate_frames[0]
            + d * cos_theta * cos_phi * local_coordinate_frames[1]
            + d * sin_theta * cos_phi * local_coordinate_frames[2]
            + d * sin_phi * local_coordinate_frames[3]
        )

        return vsite_positions

    @classmethod
    def generate_positions(
        cls,
        oe_molecule: "OEMol",
        vsite_collection: VirtualSiteCollection,
        conformer: numpy.ndarray,
    ):
        """Computes the positions of a set of virtual sites relative to a provided
        molecule in a given conformer.

        Parameters
        ----------
        oe_molecule
            The molecule to apply virtual sites to.
        vsite_collection
            The virtual site parameters to apply to the molecule
        conformer
            The conformer [A] of the molecule with shape=(n_atoms, 3) that the virtual
            sites should be placed relative to.

        Returns
        -------
            An array of virtual site positions [A] with shape=(n_vsites, 3).
        """

        vsite_parameters_by_key = {
            (parameter.smirks, parameter.type, parameter.name): parameter
            for parameter in vsite_collection.parameters
        }

        # Extract the values of the assigned parameters.
        _, assigned_parameter_map = cls._apply_virtual_sites(
            oe_molecule, vsite_collection
        )
        assigned_parameters = {
            atom_indices: [vsite_parameters_by_key[key] for key in parameter_keys]
            for atom_indices, parameter_keys in assigned_parameter_map.items()
        }

        local_frame_coordinates = numpy.vstack(
            [
                parameter.local_frame_coordinates
                for parent_indices, parameters in assigned_parameters.items()
                for parameter in parameters
            ]
        )

        # Construct the global cartesian coordinates of the v-sites.
        local_coordinate_frames = cls.build_local_coordinate_frames(
            conformer, assigned_parameters
        )
        vsite_positions = VirtualSiteGenerator.convert_local_coordinates(
            local_frame_coordinates, local_coordinate_frames, backend="numpy"
        )

        return vsite_positions


"""
Copyright 2021 The CVXPY Developers

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import warnings

import scipy  # For version checks

import cvxpy.settings as s
from cvxpy.constraints import NonNeg, Zero
from cvxpy.reductions.solution import Solution, failure_solution
from cvxpy.reductions.solvers import utilities
from cvxpy.reductions.solvers.conic_solvers.conic_solver import ConicSolver
from cvxpy.utilities.versioning import Version

class SCIPY(ConicSolver):
    """An interface for the SciPy linprog function.
    Note: This requires a version of SciPy which is >= 1.6.1
    """

    # Solver capabilities.
    MIP_CAPABLE = False
    SUPPORTED_CONSTRAINTS = ConicSolver.SUPPORTED_CONSTRAINTS

    # Map of SciPy linprog status
    STATUS_MAP = {0: s.OPTIMAL,  # Optimal
                  1: s.SOLVER_ERROR,  # Iteration limit reached
                  2: s.INFEASIBLE,  # Infeasible
                  3: s.UNBOUNDED,  # Unbounded
                  4: s.SOLVER_ERROR  # Numerical difficulties encountered
                  }

    def import_solver(self) -> None:
        """Imports the solver.
        """
        from scipy import optimize as opt
        opt  # For flake8

    def name(self):
        """The name of the solver.
        """
        return s.SCIPY

    def apply(self, problem):
        """Returns a new problem and data for inverting the new solution.

        Returns
        -------
        tuple
            (dict of arguments needed for the solver, inverse data)
        """
        data = {}
        inv_data = {self.VAR_ID: problem.x.id}

        if not problem.formatted:
            problem = self.format_constraints(problem, None)
        data[s.PARAM_PROB] = problem
        data[self.DIMS] = problem.cone_dims
        inv_data[self.DIMS] = problem.cone_dims

        constr_map = problem.constr_map
        inv_data[self.EQ_CONSTR] = constr_map[Zero]
        inv_data[self.NEQ_CONSTR] = constr_map[NonNeg]
        len_eq = problem.cone_dims.zero

        c, d, A, b = problem.apply_parameters()
        data[s.C] = c
        inv_data[s.OFFSET] = d
        data[s.A] = -A[:len_eq]
        if data[s.A].shape[0] == 0:
            data[s.A] = None
        data[s.B] = b[:len_eq].flatten()
        if data[s.B].shape[0] == 0:
            data[s.B] = None
        data[s.G] = -A[len_eq:]
        if 0 in data[s.G].shape:
            data[s.G] = None
        data[s.H] = b[len_eq:].flatten()
        if 0 in data[s.H].shape:
            data[s.H] = None
        return data, inv_data

    def solve_via_data(self, data, warm_start: bool, verbose: bool, solver_opts, solver_cache=None):
        from scipy import optimize as opt

        # Set default method which can be overriden by user inputs
        if (Version(scipy.__version__) < Version('1.6.1')):
            meth = "interior-point"
        else:
            meth = "highs"

        # Extract solver options which are not part of the options dictionary
        if solver_opts:

            # Raise error message if the parameters are not passed in
            # a dictionary called 'scipy_options'.
            if "scipy_options" not in solver_opts:
                raise ValueError("All parameters for the SCIPY solver should "
                                 "be incased within a dictionary called "
                                 "scipy_options e.g. \n"
                                 "prob.solve(solver='SCIPY', verbose=True,"
                                 " scipy_options={'method':'highs-ds', 'maxiter':10000})")

            # Raise warning if the 'method' parameter is not specified
            if "method" not in solver_opts['scipy_options']:
                warnings.warn("It is best to specify the 'method' parameter "
                              "within scipy_options. The main advantage "
                              "of this solver, is its ability to use the "
                              "HiGHS LP solvers via scipy.optimize.linprog() "
                              "which require a SciPy version >= 1.6.1 ."
                              "\n\nThe default method '{}' will be"
                              " used in this case.\n".format(meth))

            else:
                meth = solver_opts['scipy_options'].pop("method")

                # Check to see if scipy version larger than 1.6.1 is installed
                # if method chosen is one of the highs methods.
                ver = (Version(scipy.__version__) < Version('1.6.1'))
                if ((meth in ['highs-ds', 'highs-ipm', 'highs']) & ver):
                    raise ValueError("The HiGHS solvers require a SciPy version >= 1.6.1")

            # Disable the 'bounds' parameter to avoid problems with
            # canonicalised problems.
            if "bounds" in solver_opts['scipy_options']:
                raise ValueError("Please do not specify bounds through "
                                 "scipy_options. Please specify bounds "
                                 "through CVXPY.")

            # Not supported by HiGHS solvers:
            # callback = solver_opts['scipy_options'].pop("callback", None)
            # x0 = solver_opts['scipy_options'].pop("x0", None)

            # Run the optimisation using scipy.optimize.linprog
            solution = opt.linprog(data[s.C], A_ub=data[s.G], b_ub=data[s.H],
                                   A_eq=data[s.A], b_eq=data[s.B], method=meth,
                                   bounds=(None, None), options=solver_opts['scipy_options'])
        else:

            warnings.warn("It is best to specify the 'method' parameter "
                          "within scipy_options. The main advantage "
                          "of this solver, is its ability to use the "
                          "HiGHS LP solvers via scipy.optimize.linprog() "
                          "which require a SciPy version >= 1.6.1 ."
                          "\n\nThe default method '{}' will be"
                          " used in this case.\n".format(meth))

            # Run the optimisation using scipy.optimize.linprog
            solution = opt.linprog(data[s.C], A_ub=data[s.G], b_ub=data[s.H],
                                   A_eq=data[s.A], b_eq=data[s.B], method=meth,
                                   bounds=(None, None))

        if verbose is True:
            print("Solver terminated with message: " + solution.message)

        return solution

    def invert(self, solution, inverse_data):
        """Returns the solution to the original problem given the inverse_data.
        """
        status = self.STATUS_MAP[solution['status']]

        primal_vars = None
        dual_vars = None
        if status in s.SOLUTION_PRESENT:
            primal_val = solution['fun']
            opt_val = primal_val + inverse_data[s.OFFSET]
            primal_vars = {inverse_data[self.VAR_ID]: solution['x']}

            # SciPy linprog only returns duals for version >= 1.7.0
            # and method is one of 'highs', 'highs-ds' or 'highs-ipm'
            if ('ineqlin' in solution.keys()):
                eq_dual = utilities.get_dual_values(
                    -solution['eqlin']['marginals'],
                    utilities.extract_dual_value,
                    inverse_data[self.EQ_CONSTR])
                leq_dual = utilities.get_dual_values(
                    -solution['ineqlin']['marginals'],
                    utilities.extract_dual_value,
                    inverse_data[self.NEQ_CONSTR])
                eq_dual.update(leq_dual)
                dual_vars = eq_dual

            return Solution(status, opt_val, primal_vars, dual_vars, {})
        else:
            return failure_solution(status)

import scipy,from scipy,import cvxpy,from cvxpy


# THIS IS AN OLD FILE TO REIMPLEMENT, NOT MEANT TO BE IMPORTED
raise ImportError()

"""Deep Belief Netork as Bandit
"""
import copy
import cPickle
import logging
import os
import subprocess
import sys
import time
logger = logging.getLogger(__name__)

import numpy
from bson import SON, BSON

import theano
from theano import tensor

# scikit-data
from skdata.tasks import classification_train_valid_test

# XXX use scikits-learn for PCA
import pylearn_pca

from base import Bandit
from utils import json_call

try:
    RandomStreams = theano.sandbox.cuda.CURAND_RandomStreams
except:
    RandomStreams = tensor.shared_randomstreams.RandomStreams

from ht_dist2 import rSON2, one_of, rlist, uniform, lognormal, ceil_lognormal

class LogisticRegression(object):
    def __init__(self, x, w, b, params=None):
        if params is None:
            params = []
        self.input = x
        self.output = tensor.nnet.softmax(tensor.dot(x, w) + b)
        self.l1 = abs(w).sum()
        self.l2_sqr = (w**2).sum()
        self.argmax = tensor.argmax(
                tensor.dot(x, w) + b,
                axis=x.ndim - 1)
        self.w = w
        self.b = b
        self.params = params

    @classmethod
    def new(cls, input, n_in, n_out, dtype=None, name=None):
        if dtype is None:
            dtype = input.dtype
        if name is None:
            name = cls.__name__
        logger.debug('allocating params w, b: %s' % str((n_in, n_out, dtype)))
        w = theano.shared(
                numpy.zeros((n_in, n_out), dtype=dtype),
                name='%s.w' % name)
        b = theano.shared(
                numpy.zeros((n_out,), dtype=dtype),
                name='%s.b' % name)
        return cls(input, w, b, params=[w, b])

    def nll(self, target):
        """Return the negative log-likelihood of the prediction of this model under a given
        target distribution.  Passing symbolic integers here means 1-hot.
        WRITEME
        """
        return tensor.nnet.categorical_crossentropy(self.output, target)

    def errors(self, target):
        """Return a vector of 0s and 1s, with 1s on every line that was mis-classified.
        """
        if target.ndim != self.argmax.ndim:
            raise TypeError('target should have the same shape as self.argmax',
                    ('target', target.type, 'argmax', self.argmax.type))
        if target.dtype.startswith('int'):
            return theano.tensor.neq(self.argmax, target)
        else:
            raise NotImplementedError()

def sgd_updates(params, grads, stepsizes):
    """Return a list of (pairs) that can be used as updates in theano.function to implement
    stochastic gradient descent.

    :param params: variables to adjust in order to minimize some cost
    :type params: a list of variables (theano.function will require shared variables)
    :param grads: the gradient on each param (with respect to some cost)
    :type grads: list of theano expressions
    :param stepsizes: step by this amount times the negative gradient on each iteration
    :type stepsizes: [symbolic] scalar or list of one [symbolic] scalar per param
    """
    try:
        iter(stepsizes)
    except Exception:
        stepsizes = [stepsizes for p in params]
    if len(params) != len(grads):
        raise ValueError('params and grads have different lens')
    updates = [(p, p - step * gp) for (step, p, gp) in zip(stepsizes, params, grads)]
    return updates

def geom(lower, upper, round=1):
    ll = numpy.log(lower)
    lu = numpy.log(upper)
    return ceil_lognormal(.5 * (ll + lu), .4 * (lu - ll), round)

def dbn_template(dataset_name='skdata.larochelle_etal_2007.Rectangles',
        sup_min_epochs=300,
        sup_max_epochs=4000):
    template = rSON2(
        'preprocessing', one_of(
            rSON2(
                'kind', 'raw'),
            rSON2(
                'kind', 'zca',
                'energy', uniform(0.5, 1.0))),
        'dataset_name', dataset_name,
        'sup_max_epochs', sup_max_epochs,
        'sup_min_epochs', sup_min_epochs,
        'iseed', one_of(5, 6, 7, 8),
        'batchsize', one_of(20, 100),
        'lr', lognormal(numpy.log(.01), 3),
        'lr_anneal_start', geom(100, 10000),
        'l2_penalty', one_of(0, lognormal(numpy.log(1.0e-6), 2)),
        'next_layer', one_of(None,
            rSON2(
                'n_hid', geom(2**7, 2**12, round=16),
                'W_init_dist', one_of('uniform', 'normal'),
                'W_init_algo', one_of('old', 'Xavier'),
                'W_init_algo_old_multiplier', lognormal(0.0, 1.0),
                'cd_epochs', geom(1, 3000),
                'cd_batchsize', 100,
                'cd_sample_v0s', one_of(False, True),
                'cd_lr', lognormal(numpy.log(.01), 2),
                'cd_lr_anneal_start', geom(10, 10000),
                'next_layer', one_of(None,
                    rSON2(
                        'n_hid', geom(2**7, 2**12, round=16),
                        'W_init_dist', one_of('uniform', 'normal'),
                        'W_init_algo', one_of('old', 'Xavier'),
                        'W_init_algo_old_multiplier', lognormal(0.0, 1.0),
                        'cd_epochs', geom(1, 2000),
                        'cd_batchsize', 100,
                        'cd_sample_v0s', one_of(False, True),
                        'cd_lr', lognormal(numpy.log(.01), 2),
                        'cd_lr_anneal_start', geom(10, 10000),
                        'next_layer', one_of(None,
                            rSON2(
                                'n_hid', geom(2**7, 2**12, round=16),
                                'W_init_dist', one_of('uniform', 'normal'),
                                'W_init_algo', one_of('old', 'Xavier'),
                                'W_init_algo_old_multiplier', lognormal(0., 1.),
                                'cd_epochs', geom(1, 1500),
                                'cd_batchsize', 100,
                                'cd_sample_v0s', one_of(False, True),
                                'cd_lr', lognormal(numpy.log(.01), 2),
                                'cd_lr_anneal_start', geom(10, 10000),
                                'next_layer', None,
                                )))))))
    return template

def nnet1_template(dataset_name='skdata.larochelle_etal_2007.Rectangles',
            sup_min_epochs=30, # THESE ARE KINDA SMALL FOR SERIOUS RESULTS
            sup_max_epochs=400):
    template = rSON2(
        'preprocessing', one_of(
            rSON2(
                'kind', 'raw'),
            rSON2(
                'kind', 'zca',
                'energy', uniform(0.5, 1.0))),
        'dataset_name', dataset_name,
        'sup_max_epochs', sup_max_epochs,
        'sup_min_epochs', sup_min_epochs,
        'iseed', one_of(5, 6, 7, 8),
        'batchsize', one_of(20, 100),
        'lr', lognormal(numpy.log(.01), 3),
        'lr_anneal_start', geom(100, 10000),
        'l2_penalty', one_of(0, lognormal(numpy.log(1.0e-6), 3)),
        'next_layer', rSON2(
            'n_hid', geom(2**4, 2**10, round=16),
            'W_init_dist', one_of('uniform', 'normal'),
            'W_init_algo', one_of('old', 'Xavier'),
            'W_init_algo_old_multiplier', uniform(.2, 2),
            'cd_epochs', 0,
            'cd_batchsize', 100,
            'cd_sample_v0s', one_of(False, True),
            'cd_lr', lognormal(numpy.log(.01), 3),
            'cd_lr_anneal_start', geom(10, 10000),
            'next_layer', None))
    return template

def preprocess_data(config, ctrl):
    dataset = json_call(config['dataset_name'])
    train, valid, test = classification_train_valid_test(dataset)
    X_train, y_train = numpy.asarray(train[0]), numpy.asarray(train[1])
    X_valid, y_valid = numpy.asarray(valid[0]), numpy.asarray(valid[1])
    X_test, y_test = numpy.asarray(test[0]), numpy.asarray(test[1])

    if config['preprocessing']['kind'] == 'pca':
        # compute pca of input (TODO: retrieve only pca_whitened input)
        raise NotImplementedError('rewrite since cut and paste')
        (eigvals,eigvecs), centered_trainset = pylearn_pca.pca_from_examples(
                X=dataset['inputs'][:dataset['n_train']],
                max_energy_fraction=config['pca_energy'])
        eigmean = dataset['inputs'][0] - centered_trainset[0]

        whitened_inputs = pylearn_pca.pca_whiten((eigvals,eigvecs),
                dataset['inputs']-eigmean)
        ctrl.info('PCA kept %i of %i components'%(whitened_inputs.shape[1],
            dataset['n_inputs']))
    elif config['preprocessing']['kind'] == 'zca':
        (eigvals,eigvecs), centered_trainset = pylearn_pca.pca_from_examples(
                X=X_train,
                max_energy_fraction=config['preprocessing']['energy'])
        eigmean = X_train[0] - centered_trainset[0]

        def whiten(X):
            X = pylearn_pca.pca_whiten((eigvals,eigvecs),
                    X - eigmean)
            X = pylearn_pca.pca_whiten_inverse((eigvals, eigvecs),
                    X) + eigmean
            X = X.astype('float32')
            X_min = X.min()
            X_max = X.max()
            ctrl.info('ZCA min:%f max:%f' % (X_min, X_max))
            if X_min < 0 or X_max > 1.0:
                ctrl.info('ZCA clamping return value to (0, 1) interval')
                X = numpy.clip(X, 0, 1, out=X)
            return X

        X_train, X_valid, X_test = [whiten(X)
                for X in [X_train, X_valid, X_test]]

    elif config['preprocessing']['kind'] == 'normalize':
        raise NotImplementedError('rewrite since cut and paste')
        n_train=dataset['n_train']
        whitened_inputs = dataset['inputs']
        whitened_inputs = whitened_inputs - whitened_inputs[:n_train].mean(axis=0)
        whitened_inputs /= whitened_inputs[:n_train].std(axis=0)+1e-7
    elif config['preprocessing']['kind'] == 'raw':
        pass
    else:
        raise ValueError(
                'unrecognized preprocessing',
                config['preprocessing']['kind'])

    for Xy in 'X', 'y':
        for suffix in 'train', 'valid', 'test':
            varname = '%s_%s'%(Xy, suffix)
            var = locals()[varname]
            ctrl.info('%s shape=%s max=%f min=%f' % (
                varname,
                var.shape,
                var.max(),
                var.min()))

    s_X_train = theano.shared(X_train)
    s_y_train = theano.shared(y_train)
    s_X_valid = theano.shared(X_valid)
    s_y_valid = theano.shared(y_valid)
    s_X_test = theano.shared(X_test)
    s_y_test = theano.shared(y_test)

    return (dataset,
            (s_X_train, s_y_train),
            (s_X_valid, s_y_valid),
            (s_X_test, s_y_test))

def train_rbm(s_rng, s_idx, s_batchsize, s_features, W, vbias, hbias, n_in,
        n_hid, batchsize, sample_v0s,
        cdlr, n_epochs, n_batches_per_epoch, lr_anneal_start,
        givens={},
        time_limit=None):
    logger.info('rbm training n_in=%i n_hid=%i batchsize=%i' % (
        n_in, n_hid, batchsize))
    v0m = s_features
    if sample_v0s:
        v0s = tensor.cast(
                s_rng.uniform(size=(batchsize, n_in)) < v0m,
                'float32')
    else:
        v0s = v0m

    h0m = tensor.nnet.sigmoid(tensor.dot(v0s, W) + hbias)
    h0s = tensor.cast(s_rng.uniform(size=(batchsize, n_hid)) < h0m, 'float32')
    v1m = tensor.nnet.sigmoid(tensor.dot(h0s, W.T)+vbias)
    v1s = tensor.cast(s_rng.uniform(size=(batchsize, n_in)) < v1m, 'float32')
    h1m = tensor.nnet.sigmoid(tensor.dot(v1s, W) + hbias)

    s_lr = tensor.scalar(dtype='float32')

    logger.debug('compiling cd1_fn')
    cd1_fn = theano.function([s_idx, s_batchsize, s_lr],
            [abs(v0m-v1m).mean()],
            updates={
                W: W + s_lr * (
                    tensor.dot(v0s.T, h0m) - tensor.dot(v1s.T, h1m)),
                vbias: vbias + s_lr * (
                    (v0s - v1s).sum(axis=0)),
                hbias: hbias + s_lr * (
                    (h0m - h1m).sum(axis=0)),
                },
            givens=givens)
    for epoch in xrange(n_epochs):
        costs = []
        if time_limit and time.time() > time_limit:
            break
        e_lr = cdlr * min(1, (float(lr_anneal_start)/(epoch+1)))
        for batch_idx in xrange(n_batches_per_epoch):
            costs.append(cd1_fn(batch_idx, batchsize, e_lr))
        if not epoch % 10:
            logger.info('CD1 epoch:%i  avg L1: %f'% (epoch, numpy.mean(costs)))
    if costs:
        return dict(final_recon_l1=float(numpy.mean(costs)),)
    else:
        return dict(final_recon_l1=float('nan'))

_dataset_cache = {}

class DBN_Base(Bandit):
    def dryrun_config(self, *args, **kwargs):
        return dict(
                lr=.01,
                sup_max_epochs=500,
                sup_min_epochs=50,
                batchsize=10,
                preprocessing=dict(kind='zca', energy=0.8),
                iseed=5,
                n_layers=1,
                next_layer = dict(
                    n_hid=50,
                    W_init_dist='uniform',
                    W_init_algo='Xavier',
                    cd_epochs=100,
                    cd_batchsize=50,
                    cd_sample_v0s=True,
                    cd_lr=0.1,
                    cd_lr_anneal_start=3,
                    next_layer = dict(
                        n_hid=75,
                        W_init_dist='uniform',
                        W_init_algo='old',
                        W_init_algo_old_multiplier=2.2,
                        cd_epochs=70,
                        cd_batchsize=10,
                        cd_sample_v0s=False,
                        cd_lr=0.01,
                        cd_lr_anneal_start=30
                        ),
                    ),
                l2_penalty=0.1,
                lr_anneal_start=20,
                dataset_name='skdata.larochelle_etal_2007.Rectangles',
                )

    @classmethod
    def evaluate(cls, config, ctrl):
        time_limit = time.time() + 60 * 60 # 1hr from now
        rval = SON(dbn_train_fn_version=1)

        ctrl.info('starting dbn_train_fn')
        kv = config.items()
        kv.sort()
        for k,v in kv:
            ctrl.info('key=%s\t%s' %(k,str(v)))

        rng = numpy.random.RandomState(config['iseed'])
        s_rng = RandomStreams(int(rng.randint(2**30)))

        dataset, train_Xy, valid_Xy, test_Xy = preprocess_data(config, ctrl)

        # allocate learning function parameters
        s_inputs_all = tensor.fmatrix('inputs')
        s_labels_all = tensor.ivector('labels')
        s_idx = tensor.lscalar('batch_idx')
        s_batchsize=tensor.lscalar('batch_size')
        s_low = s_idx * s_batchsize
        s_high = s_low + s_batchsize
        s_inputs = s_inputs_all[s_low:s_high]
        s_labels = s_labels_all[s_low:s_high]
        s_lr = tensor.scalar('lr')
        s_features = s_inputs # s_features will be modified in the model-building loop

        weights = []
        vbiases = []
        hbiases = []

        n_inputs_i = valid_Xy[0].get_value(borrow=True).shape[1]

        rval['cd_reports'] = []

        try:
            layer_config = config['next_layer']
            # allocate model parameters
            while layer_config:
                i = len(rval['cd_reports'])
                n_hid_i = layer_config['n_hid']
                if layer_config['W_init_dist']=='uniform':
                    W = rng.uniform(low=-1,high=1,size=(n_hid_i, n_inputs_i)).T.astype('float32')
                elif layer_config['W_init_dist'] == 'normal':
                    W = rng.randn(n_hid_i, n_inputs_i).T.astype('float32')
                else:
                    raise ValueError('W_init_dist', layer_config['W_init_dist'])

                if layer_config['W_init_algo'] == 'old':
                    #N.B. the weights are transposed so that as the number of hidden units changes,
                    # the first hidden units are always the same vectors.
                    # this makes it easier to isolate the effect of random initialization
                    # from the other hyper-parameters under review
                    W *= layer_config['W_init_algo_old_multiplier'] / numpy.sqrt(n_inputs_i)
                elif layer_config['W_init_algo'] == 'Xavier':
                    W *= numpy.sqrt(6.0 / (n_inputs_i + n_hid_i))
                else:
                    raise ValueError(layer_config['W_init_algo'])

                layer_idx = len(rval['cd_reports'])
                weights.append(theano.shared(W, 'W_%i' % layer_idx))
                hbiases.append(theano.shared(numpy.zeros(n_hid_i, dtype='float32'),
                    'h_%i' % layer_idx))
                vbiases.append(theano.shared(numpy.zeros(n_inputs_i, dtype='float32'),
                    'v_%i' % layer_idx))
                del W

                # allocate RBM training function for this layer
                # this version re-calculates the training set every time
                # TODO: cache the training set for each layer
                # TODO: consider sparsity?
                # TODO: consider momentum?
                if layer_config['cd_epochs']:
                    cd_report = train_rbm(
                            s_rng, s_idx, s_batchsize, s_features,
                            W=weights[-1],
                            vbias=vbiases[-1],
                            hbias=hbiases[-1],
                            n_in=n_inputs_i,
                            n_hid=n_hid_i,
                            batchsize=layer_config['cd_batchsize'],
                            sample_v0s=layer_config['cd_sample_v0s'],
                            cdlr=layer_config['cd_lr'] / float(layer_config['cd_batchsize']),
                            n_epochs=layer_config['cd_epochs'],
                            n_batches_per_epoch=dataset.descr['n_train'] // layer_config['cd_batchsize'],
                            lr_anneal_start=layer_config['cd_lr_anneal_start'],
                            givens = {
                                s_inputs_all: tensor.as_tensor_variable(train_Xy[0])
                                },
                            time_limit=time_limit
                            )
                else:
                    cd_report = None
                rval['cd_reports'].append(cd_report)

                # update s_features to point to top layer
                s_features = tensor.nnet.sigmoid(
                        tensor.dot(s_features, weights[-1]) + hbiases[-1])
                n_inputs_i = n_hid_i
                layer_config = layer_config.get('next_layer', None)

        except (MemoryError,):
            rval['abort'] = 'MemoryError'
            rval['status'] = 'ok'
            rval['loss'] = 1.0
            rval['best_epoch_valid'] = 0.0
            return rval

        # allocate model

        logreg = LogisticRegression.new(s_features, n_in=n_inputs_i,
                n_out=dataset.descr['n_classes'])
        traincost = logreg.nll(s_labels).mean()
        def ssq(X):
            return (X**2).sum()
        traincost = traincost + config['l2_penalty'] * (
                sum([ssq(w_i) for w_i in weights]) + ssq(logreg.w))
        # params = weights+hbiases+vbiases+logreg.params
        # vbiases are not involved in the supervised network
        params = weights + hbiases + logreg.params
        train_logreg_fn = theano.function([s_idx, s_lr],
                [logreg.nll(s_labels).mean()],
                updates=sgd_updates(
                    params=params,
                    grads=tensor.grad(traincost, params),
                    stepsizes=[s_lr] * len(params)),
                givens={s_batchsize:config['batchsize'],
                    s_inputs_all: tensor.as_tensor_variable(train_Xy[0]),
                    s_labels_all: train_Xy[1]})
        valid_logreg_fn = theano.function([s_idx],
            logreg.errors(s_labels).mean(),
            givens={s_batchsize:config['batchsize'],
                s_inputs_all: tensor.as_tensor_variable(valid_Xy[0]),
                s_labels_all: valid_Xy[1]})
        test_logreg_fn = theano.function([s_idx],
            logreg.errors(s_labels).mean(),
            givens={s_batchsize:config['batchsize'],
                s_inputs_all: tensor.as_tensor_variable(test_Xy[0]),
                s_labels_all: test_Xy[1]})

        rval['best_epoch'] = -1
        rval['best_epoch_valid'] = -1
        rval['best_epoch_train'] = -1
        rval['best_epoch_test'] = -1
        rval['status'] = 'ok'
        valid_rate=-1
        test_rate=-1
        train_rate=-1

        n_train_batches = dataset.descr['n_train'] // config['batchsize']
        n_valid_batches = dataset.descr['n_valid'] // config['batchsize']
        n_test_batches = dataset.descr['n_test'] // config['batchsize']

        n_iters = 0
        for epoch in xrange(config['sup_max_epochs']):
            e_lr = config['lr']
            e_lr *= min(1, config['lr_anneal_start'] / float(n_iters+1)) #anneal learning rate
            valid_rate = float(1 - numpy.mean([valid_logreg_fn(i)
                for i in range(n_valid_batches)]))
            valid_rate_std_thresh = 0.5 * numpy.sqrt(valid_rate *
                    (1 - valid_rate) / (n_valid_batches * config['batchsize']))

            if valid_rate > (rval['best_epoch_valid']+valid_rate_std_thresh):
                rval['best_epoch'] = epoch
                rval['best_epoch_test'] = test_rate
                rval['best_epoch_valid'] = valid_rate
                rval['best_epoch_train'] = train_rate
                best_params = copy.deepcopy(params)
            logger.info('Epoch=%i best epoch %i valid %f test %f best_epoch_train %f prev_train %f'%(
                epoch, rval['best_epoch'], rval['best_epoch_valid'], rval['best_epoch_test'],
                    rval['best_epoch_train'], train_rate))
            #ctrl.info('Epoch %i train nll: %f'%(epoch, train_rate))
            ctrl.checkpoint(rval)

            if epoch > config['sup_min_epochs'] and epoch > 2*rval['best_epoch']:
                break
            if time.time() > time_limit:
                break
            train_rate = float(numpy.mean([train_logreg_fn(i,e_lr) for i in
                range(n_train_batches)]))
            if not numpy.isfinite(train_rate):
                do_test = False
                rval['status'] = 'fail'
                rval['status_info'] = 'train_rate %f' % train_rate
                break
            ++n_iters

        do_test = 1
        if do_test and rval['status'] == 'ok':
            # copy best params back into place
            for p, bp in zip(params, best_params):
                p.set_value(bp.get_value())
            rval['best_epoch_test'] = 1 - float(
                    numpy.mean(
                        [test_logreg_fn(i) for i in range(n_test_batches)]))
            rval['loss'] = 1.0 - rval['best_epoch_valid']
        ctrl.info('rval: %s' % str(rval))
        return rval

    @classmethod
    def loss(cls, result, config=None):
        """Extract the scalar-valued loss from a result document
        """
        try:
            if numpy.isnan(float(result['loss'])):
                return None
            else:
                return float(result['loss'])
        except KeyError, TypeError:
            return None

    @classmethod
    def loss_variance(cls, result, config=None):
        if config['dataset_name'] not in _dataset_cache:
            _dataset_cache[config['dataset_name']] = json_call(
                config['dataset_name'])
        dataset = _dataset_cache[config['dataset_name']]
        n_valid = dataset.descr['n_valid']
        p = cls.loss(result, config)
        if p is None:
            return None
        return p * (1.0 - p) / (n_valid - 1)

    @classmethod
    def true_loss(cls, result, config=None):
        try:
            rval =  float(1 - result['best_epoch_test'])
            if 0 <= rval <= 1:
                return rval
        except (KeyError, TypeError):
            return None

    @classmethod
    def status(cls, result, config=None):
        """Extract the job status from a result document
        """
        if (result['status'] == 'ok' and
            (cls.loss(result) is None
                or cls.true_loss(result) is None)):
            return 'fail'
        else:
            return result['status']

def DBN_Convex():
    return DBN_Base(
            dbn_template(
                dataset_name='skdata.larochelle_etal_2007.Convex'))

def DBN_MRBI():
    ds =  'skdata.larochelle_etal_2007.MNIST_RotatedBackgroundImages'
    return DBN_Base(dbn_template(dataset_name=ds))

class Dummy_DBN_Base(Bandit):
    """
    A DBN_Base stub.

    This class is used in unittests of optimization algorithms to ensure they
    can deal with large nested specifications that include lots of distribution
    types.

    The evaluate function simply returns a random score.
    """
    def __init__(self):
        Bandit.__init__(self, template=dbn_template())
        self.rng = numpy.random.RandomState(234)

    def evaluate(self, argd, ctrl):
        rval = dict(dbn_train_fn_version=-1)
        # XXX: TODO: make up a loss function that depends on argd.
        rval['status'] = 'ok'
        rval['best_epoch_valid'] = float(self.rng.rand())
        rval['loss'] = 1.0 - rval['best_epoch_valid']
        return rval

    def loss_variance(self, result, config=None):
        return 0.01**2


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

import re
import numpy as np

from .dbmgr import get_cursor

class Element(object):
	"""Elemental data

	...

	Parameters
	----------
	x : type
		Description of parameter `x`.

	Attributes
	----------

	Methods
	-------

	Notes
	-----

	References
	----------

	Examples
	--------

	"""

	def __init__(self, element):
		self.element = element.title()
		self._meta = None

	@property
	def meta(self):
		if self._meta is None:
			from scipy.interpolate import interp1d

			self._meta = {}
			zg = get_cursor('ziegler')
			self._meta['Z'] = [str(i[2]).split(':')[0] for i in zg.execute('SELECT * FROM compounds WHERE compound=?',(self.element,))][0]
			self._meta['mass'], self._meta['density'] = [(i[1],i[2]) for i in zg.execute('SELECT * FROM weights WHERE Z=?',(self._meta['Z'],))][0]
			coeff = np.array([i[1:] for i in zg.execute('SELECT * FROM mass_coeff WHERE Z=? ORDER BY energy',(self._meta['Z'],))])
			self._meta['mass_coeff'] = interp1d(coeff[:,0], coeff[:,1], bounds_error=False, fill_value=0.0)
			self._meta['mass_coeff_en'] = interp1d(coeff[:,0], coeff[:,2], bounds_error=False, fill_value=0.0)
			db = get_cursor('decay')
			self._meta['abundances'] = {str(i[1]):i[12] for i in db.execute('SELECT * FROM chart WHERE element=? AND abundance>0.0',(self.element,))}
			self._meta['isotopes'] = sorted([i for i in self._meta['abundances']])

		return self._meta

	@property
	def Z(self):
		return self.meta['Z']

	@property
	def mass(self):
		return self.meta['mass']

	@property
	def isotopes(self):
		return self.meta['isotopes']

	@property
	def abundances(self):
		return self.meta['abundances']

	@property
	def mass_coeff(self):
		return self.meta['mass_coeff']

	@property
	def mass_coeff_en(self):
		return self.meta['mass_coeff_en']

	def attenuation(self, E, x=1.0):
		return np.exp(-self.mass_coeff(E)*self.density*x)

	def transmission(self, E, x=1.0):
		return 1.0-np.exp(-self.mass_coeff(E)*self.density*x)

	@property
	def density(self):
		return self.meta['density']

class Isotope(object):
	"""Isotopic data

	...

	Parameters
	----------
	x : type
		Description of parameter `x`.

	Attributes
	----------

	Methods
	-------

	Notes
	-----

	References
	----------

	Examples
	--------

	"""

	def __init__(self, istp):
		if istp=='1n' or istp=='1ng':
			self.element, self.A, self.isomer = 'n', 1, 'g'
		else:
			self.element = ''.join(re.findall('[A-Z]+',istp))
			if istp.startswith('nat'):
				self.A = 'nat'
			else:
				self.A = int(istp.split(self.element)[0])
			self.isomer = istp.split(self.element)[1]
		self.isotope = str(self.A)+self.element
		if self.isomer=='' and type(self.A)==int:
			self.isomer = 'g'
		if self.isomer=='m':
			self.isomer = 'm1'
		self.name = self.isotope+self.isomer
		self._meta = None

	@property
	def meta(self):
		if self._meta is None:
			self._meta = {}
			self.db = get_cursor('decay')
			for i in ['SFY','gm','el','bm','bp','al']:
				self._meta[i] = None
			Q = list(self.db.execute('SELECT * FROM chart WHERE isotope=?',(self.isotope,)))
			q = [i for i in Q if i[3]==self.isomer]
			if len(q):
				q = q[0]
			elif len(Q):
				q = Q[0]
			else:
				q = [0,0,0,0,'',0,None,None,None,None,None,None,None,None,None,None,'']
			self._meta['E_level'] = q[2]
			self._meta['J_pi'] = str(q[4]) if q[4] else '?'
			self._meta['Z'] = q[6]
			self._meta['N'] = q[7]
			self._meta['stable'] = bool(q[9])
			self._meta['t_half'] = q[10]
			self._meta['unc_t_half'] = q[11]
			self._meta['abundance'] = q[12]
			self._meta['unc_abundance'] = q[13]
			if self._meta['stable'] and self._meta['abundance'] is None:
				self._meta['abundance'] = 100.0
				self._meta['unc_abundance'] = 0.0
			self._meta['mass'] = q[14]
			self._meta['Delta'] = q[15]
			self._meta['decay_mode'] = list(map(lambda i:[float(p) if n==2 else p for n,p in enumerate(i.split(':'))], str(q[16]).split(',')))
			state = '' if len(Q)==1 else (self.isomer[0] if len(Q)==2 else self.isomer)
			self._meta['TeX'] = r'$^{'+str(self.A)+state+r'}$'+self.element.title()
		return self._meta

	@property
	def E_level(self):
		return self.meta['E_level']

	@property
	def J_pi(self):
		return self.meta['J_pi']

	@property
	def Z(self):
		return self.meta['Z']

	@property
	def N(self):
		return self.meta['N']

	@property
	def stable(self):
		return self.meta['stable']

	@property
	def mass(self):
		return self.meta['mass']

	@property
	def Delta(self):
		return self.meta['Delta']

	@property
	def TeX(self):
		return self.meta['TeX']

	def __str__(self):
		return self.name

	def half_life(self, units='s', unc=False):
		"""Description

		...

		Parameters
		----------
		x : type
			Description of parameter `x`.

		Returns
		-------

		Notes
		-----

		References
		----------

		Examples
		--------

		"""

		if self.stable:
			return (np.inf,0.0) if unc else np.inf
		half_conv = {'ns':1e-9,'us':1e-6,'ms':1e-3,'s':1.0,'m':60.0,'h':3600.0,'d':86400.0,'y':31557.6E3,'ky':31557.6E6}[units]
		if unc:
			return self.meta['t_half']/half_conv,self.meta['unc_t_half']/half_conv
		return self.meta['t_half']/half_conv

	def decay_const(self, units='s', unc=False):
		"""Description

		...

		Parameters
		----------
		x : type
			Description of parameter `x`.

		Returns
		-------

		Notes
		-----

		References
		----------

		Examples
		--------

		"""

		if self.stable:
			return (0.0,0.0) if unc else 0.0
		if unc:
			T2,uT2 = self.half_life(units, True)
			return np.log(2.0)/T2,np.log(2.0)*uT2/T2**2
		return np.log(2.0)/self.half_life(units)

	def optimum_units(self):
		"""Description

		...

		Parameters
		----------
		x : type
			Description of parameter `x`.

		Returns
		-------

		Notes
		-----

		References
		----------

		Examples
		--------

		"""

		opt = ['ns']
		for units in ['us','ms','s','m','h','d','y']:
			if self.half_life(units)>1.0:
				opt.append(units)
		return opt[-1]

	def abundance(self, unc=False):
		"""Description

		...

		Parameters
		----------
		x : type
			Description of parameter `x`.

		Returns
		-------

		Notes
		-----

		References
		----------

		Examples
		--------

		"""

		if unc:
			return self.meta['abundance'], self.meta['unc_abundance']
		return self.meta['abundance']

	def get_SFY(self, unc=False, closest_SFY=False):
		"""Description

		...

		Parameters
		----------
		x : type
			Description of parameter `x`.

		Returns
		-------

		Notes
		-----

		References
		----------

		Examples
		--------

		"""

		if self.meta['SFY'] is None:
			self.meta['SFY'] = [(str(i[1]),i[2],i[3]) for i in self.db.execute('SELECT * FROM SFY WHERE parent=?',(self.name,))]
		SFY = list(self.meta['SFY'])
		if len(SFY)==0 and closest_SFY:
			itps = list(set([str(i[0]) for i in self.db.execute('SELECT parent FROM SFY')]))
			dA = [abs(self.A-int(i[:3])) for i in itps]
			SFY = [(str(i[1]),i[2],i[3]) for i in self.db.execute('SELECT * FROM SFY WHERE parent=?',(itps[dA.index(min(dA))],))]
		if unc:
			return [[i[0], i[1], i[2]] for i in SFY]
		return [[i[0], i[1]] for i in SFY]

	def decay_products(self, closest_SFY=False):
		"""Description

		...

		Parameters
		----------
		x : type
			Description of parameter `x`.

		Returns
		-------

		Notes
		-----

		References
		----------

		Examples
		--------

		"""

		prods = []
		for (mode, product, br) in self.meta['decay_mode']:
			if product=='SFY':
				prods += [[i, br*y] for i,y in self.get_SFY(False, closest_SFY)]
			else:
				prods.append([product, br])
		return [[i, br] for i,br in prods if br>1E-8]

	def gammas(self, I_lim=[None,None], E_lim=[None,None], xrays=False, dE_511=3.5):
		"""Description

		...

		Parameters
		----------
		x : type
			Description of parameter `x`.

		Returns
		-------

		Notes
		-----

		References
		----------

		Examples
		--------

		"""

		if self.meta['gm'] is None:
			self.meta['gm'] = [[float(i[3]),float(i[4]),float(i[5]),str(i[6])] for i in self.db.execute('SELECT * FROM gammas WHERE isotope=? AND isomer=?',(self.isotope,self.isomer))]
		gammas = list(self.meta['gm'])
		for n,L in enumerate([E_lim,I_lim]):
			if L[0] is not None:
				gammas = [g for g in gammas if g[n]>=L[0]]
			if L[1] is not None:
				gammas = [g for g in gammas if g[n]<=L[1]]
		if not xrays:
			gammas = [g for g in gammas if g[3]=='']
		return {l:[g[n] for g in gammas if abs(g[0]-511.0)>=dE_511] for n,l in enumerate(['E','I','dI','notes'])}

	def electrons(self,I_lim=(None,None),E_lim=(None,None),CE_only=False,Auger_only=False):
		"""Description

		...

		Parameters
		----------
		x : type
			Description of parameter `x`.

		Returns
		-------

		Notes
		-----

		References
		----------

		Examples
		--------

		"""

		if self.meta['el'] is None:
			self.meta['el'] = [[float(i[3]),float(i[4]),float(i[5]),str(i[6])] for i in self.db.execute('SELECT * FROM electrons WHERE isotope=? AND isomer=?',(self.isotope,self.isomer))]
		electrons = list(self.meta['el'])
		for n,L in enumerate([E_lim,I_lim]):
			if L[0] is not None:
				electrons = [e for e in electrons if e[n]>=L[0]]
			if L[1] is not None:
				electrons = [e for e in electrons if e[n]<=L[1]]
		if CE_only:
			electrons = [e for e in electrons if e[3].startswith('CE')]
		if Auger_only:
			electrons = [e for e in electrons if e[3].startswith('Aug')]
		return {l:[e[n]for e in electrons] for n,l in enumerate(['E','I','dI','notes'])}

	def beta_minus(self,I_lim=(None,None),Endpoint_lim=(None,None)):
		"""Description

		...

		Parameters
		----------
		x : type
			Description of parameter `x`.

		Returns
		-------

		Notes
		-----

		References
		----------

		Examples
		--------

		"""

		if self.meta['bm'] is None:
			self.meta['bm'] = [[float(i[3]),float(i[4]),float(i[5]),float(i[6])] for i in self.db.execute('SELECT * FROM beta_minus WHERE isotope=? AND isomer=?',(self.isotope,self.isomer))]
		betas = list(self.meta['bm'])
		for n,L in zip([3,1],[Endpoint_lim,I_lim]):
			if L[0] is not None:
				betas = [b for b in betas if b[n]>=L[0]]
			if L[1] is not None:
				betas = [b for b in betas if b[n]<=L[1]]
		return {l:[b[n] for b in betas] for n,l in enumerate(['muE','I','dI','endE'])}

	def beta_plus(self, I_lim=(None,None), Endpoint_lim=(None,None)):
		"""Description

		...

		Parameters
		----------
		x : type
			Description of parameter `x`.

		Returns
		-------

		Notes
		-----

		References
		----------

		Examples
		--------

		"""

		if self.meta['bp'] is None:
			self.meta['bp'] = [[float(i[3]),float(i[4]),float(i[5]),float(i[6])] for i in self.db.execute('SELECT * FROM beta_plus WHERE isotope=? AND isomer=?',(self.isotope,self.isomer))]
		betas = list(self.meta['bp'])
		for n,L in zip([3,1],[Endpoint_lim,I_lim]):
			if L[0] is not None:
				betas = [b for b in betas if b[n]>=L[0]]
			if L[1] is not None:
				betas = [b for b in betas if b[n]<=L[1]]
		return {l:[b[n] for b in betas] for n,l in enumerate(['muE','I','dI','endE'])}

	def alphas(self, I_lim=(None,None), E_lim=(None,None)):
		"""Description

		...

		Parameters
		----------
		x : type
			Description of parameter `x`.

		Returns
		-------

		Notes
		-----

		References
		----------

		Examples
		--------

		"""

		if self.meta['al'] is None:
			self.meta['al'] = [[float(i[3]),float(i[4]),float(i[5])] for i in self.db.execute('SELECT * FROM alphas WHERE isotope=? AND isomer=?',(self.isotope,self.isomer))]
		alphas = list(self.meta['al'])
		for n,L in enumerate([E_lim,I_lim]):
			if L[0] is not None:
				alphas = [a for a in alphas if a[n]>=L[0]]
			if L[1] is not None:
				alphas = [a for a in alphas if a[n]<=L[1]]
		return {l:[a[n] for a in alphas] for n,l in enumerate(['E','I','dI'])}

	def dose_rate(self, activity=1.0, distance=30.0, units='R/hr'):
		"""Description

		...

		Parameters
		----------
		x : type
			Description of parameter `x`.

		Returns
		-------

		Notes
		-----

		References
		----------

		Examples
		--------

		"""

		def beta2(E_MeV, m_amu):
			return 1.0-(1.0/(1.0+(E_MeV/m_amu))**2)

		def e_range(E_keV):
			dEdx = lambda b2, t: 0.17*((np.log(3.61E5*t*np.sqrt(t+2)))+(0.5*(1-b2)*(1+t**2/8-(2*t+1)*np.log(2)))-4.312)/b2
			E = np.linspace(1E-12, E_keV*1E-3, 100)
			return np.trapz(1.0/((1.0+(E*7.22/800.0))*dEdx(beta2(E, 0.5109), E/0.5109)), E)

		def pos_range(E_keV):
			dEdx = lambda b2, t: 0.17*((np.log(3.61E5*t*np.sqrt(t+2)))+(np.log(2)-(b2/24)*(23+14/(t+2)+10/(t+2)**2+4/(t+2)**3))-4.312)/b2
			E = np.linspace(1E-12, E_keV*1E-3, 100)
			return np.trapz(1.0/((1.0+(E*7.22/800.0))*dEdx(beta2(E, 0.5109), E/0.5109)), E)

		def alpha_range(E_keV):
			dEdx = lambda b2: 0.17*4.0*((np.log(1.02E6*b2/(1.0-b2))-b2)-4.312)/b2
			return np.trapz(1.0/dEdx(beta2(np.linspace(1E-9, E_keV*1E-3, 100), 3.7284E3)), np.linspace(1E-9, E_keV*1E-3, 100))

		dose = {}
		gm = self.gammas(xrays=True, dE_511=0.0)
		al = self.alphas()
		bm = self.beta_minus()
		bp = self.beta_plus()
		el = self.electrons()

		dose['gammas'] = 1.4042E-12*np.sum(np.array(gm['E'])*np.array(gm['I']))*activity/distance**2
		dose['alphas'] = 5.2087E-14*np.sum([al['I'][n]*e/alpha_range(e) for n,e in enumerate(al['E'])])*activity/distance**2
		dose['beta_minus'] = 5.2087E-14*np.sum([bm['I'][n]*e/e_range(e) for n,e in enumerate(bm['muE'])])*activity/distance**2
		dose['beta_plus'] = 5.2087E-14*np.sum([bp['I'][n]*e/pos_range(e) for n,e in enumerate(bp['muE'])])*activity/distance**2
		dose['electrons'] = 5.2087E-14*np.sum([el['I'][n]*e/e_range(e) for n,e in enumerate(el['E'])])*activity/distance**2
		return dose