# Very slightly adapted from perses https://github.com/choderalab/perses
# License: MIT
# OpenFE note: eventually we aim to move this to openmmtools where possible

import numpy as np
import warnings
import copy
from openmmtools.alchemy import AlchemicalState

class LambdaProtocol(object):
    """Protocols for perturbing each of the compent energy terms in alchemical
    free energy simulations.

    TODO
    ----
    * Class needs cleaning up and made more consistent
    """

    default_functions = {'lambda_sterics_core':
                         lambda x: x,
                         'lambda_electrostatics_core':
                         lambda x: x,
                         'lambda_sterics_insert':
                         lambda x: 2.0 * x if x < 0.5 else 1.0,
                         'lambda_sterics_delete':
                         lambda x: 0.0 if x < 0.5 else 2.0 * (x - 0.5),
                         'lambda_electrostatics_insert':
                         lambda x: 0.0 if x < 0.5 else 2.0 * (x - 0.5),
                         'lambda_electrostatics_delete':
                         lambda x: 2.0 * x if x < 0.5 else 1.0,
                         'lambda_bonds':
                         lambda x: x,
                         'lambda_angles':
                         lambda x: x,
                         'lambda_torsions':
                         lambda x: x
                         }

    # lambda components for each component,
    # all run from 0 -> 1 following master lambda
    def __init__(self, functions='default', windows=10, lambda_schedule=None):
        """Instantiates lambda protocol to be used in a free energy
        calculation. Can either be user defined, by passing in a dict, or using
        one of the pregenerated sets by passing in a string 'default', 'namd'
        or 'quarters'

        All protocols must begin and end at 0 and 1 respectively. Any energy
        term not defined in `functions` dict will be set to the function in
        `default_functions`

        Pre-coded options:
        default : ele and LJ terms of the old system are turned off between
        0.0 -> 0.5 ele and LJ terms of the new system are turned on between
        0.5 -> 1.0 core terms treated linearly

        quarters : 0.25 of the protocol is used in turn to individually change
        the (a) off old ele, (b) off old sterics, (c) on new sterics (d) on new
        ele core terms treated linearly

        namd : follows the protocol outlined here:
        https://pubs.acs.org/doi/full/10.1021/acs.jcim.9b00362#
        Jiang, Wei, Christophe Chipot, and Benoît Roux. "Computing Relative
        Binding Affinity of Ligands to Receptor: An Effective Hybrid
        Single-Dual-Topology Free-Energy Perturbation Approach in NAMD."
        Journal of chemical information and modeling 59.9 (2019): 3794-3802.

        ele-scaled : all terms are treated as in default, except for the old
        and new ele these are scaled with lambda^0.5, so as to be linear in
        energy, rather than lambda

        Parameters
        ----------
        functions : str or dict
            One of the predefined lambda protocols
            ['default','namd','quarters'] or a dictionary. Default "default".
        windows : int
            Number of windows which this lambda schedule is intended to be used
            with. This value is used to validate the lambda function.
        lambda_schedule : list of floats
            Schedule of lambda windows to be sampled. If ``None`` will default
            to a linear spacing of windows as defined by
            ``np.linspace(0. ,1. ,windows)``. Default ``None``.

        Attributes
        ----------
        functions : dict
            Lambda protocol to be used.
        lambda_schedule : list
            Schedule of windows to be sampled.
        """
        self.functions = copy.deepcopy(functions)

        # set the lambda schedule
        self.lambda_schedule = self._validate_schedule(lambda_schedule,
                                                       windows)
        if lambda_schedule:
            self.lambda_schedule = lambda_schedule
        else:
            self.lambda_schedule = np.linspace(0., 1., windows)

        if type(self.functions) == dict:
            self.type = 'user-defined'
        elif type(self.functions) == str:
            self.functions = None  # will be set later
            self.type = functions

        if self.functions is None:
            if self.type == 'default':
                self.functions = copy.deepcopy(
                                     LambdaProtocol.default_functions)
            elif self.type == 'namd':
                self.functions = {
                    'lambda_sterics_core': lambda x: x,
                    'lambda_electrostatics_core': lambda x: x,
                    'lambda_sterics_insert': lambda x: (3. / 2.) * x if x < (2. / 3.) else 1.0,
                    'lambda_sterics_delete': lambda x: 0.0 if x < (1. / 3.) else (x - (1. / 3.)) * (3. / 2.),
                    'lambda_electrostatics_insert': lambda x: 0.0 if x < 0.5 else 2.0 * (x - 0.5),
                    'lambda_electrostatics_delete': lambda x: 2.0 * x if x < 0.5 else 1.0,
                    'lambda_bonds': lambda x: x,
                    'lambda_angles': lambda x: x,
                    'lambda_torsions': lambda x: x
                }
            elif self.type == 'quarters':
                self.functions = {
                    'lambda_sterics_core': lambda x: x,
                    'lambda_electrostatics_core': lambda x: x,
                    'lambda_sterics_insert': lambda x: 0. if x < 0.5 else 1 if x > 0.75 else 4 * (x - 0.5),
                    'lambda_sterics_delete': lambda x: 0. if x < 0.25 else 1 if x > 0.5 else 4 * (x - 0.25),
                    'lambda_electrostatics_insert': lambda x: 0. if x < 0.75 else 4 * (x - 0.75),
                    'lambda_electrostatics_delete': lambda x: 4.0 * x if x < 0.25 else 1.0,
                    'lambda_bonds': lambda x: x,
                    'lambda_angles': lambda x: x,
                    'lambda_torsions': lambda x: x
                }
            elif self.type == 'ele-scaled':
                self.functions = {
                    'lambda_electrostatics_insert': lambda x: 0.0 if x < 0.5 else ((2*(x-0.5))**0.5),
                    'lambda_electrostatics_delete': lambda x: (2*x)**2 if x < 0.5 else 1.0
                }
            elif self.type == 'user-defined':
                self.functions = functions
            else:
                errmsg = f"LambdaProtocol type : {self.type} not recognised "
                raise ValueError(errmsg)

        self._validate_functions(n=windows)
        self._check_for_naked_charges()

    @staticmethod
    def _validate_schedule(schedule, windows):
        """
        Checks that the input lambda schedule is valid.

        Rules are:
          - Must begin at 0 and end at 1
          - Must be monotonically increasing

        Parameters
        ----------
        schedule : list of floats
            The lambda schedule. If ``None`` the method returns
            ``np.linspace(0. ,1. ,windows)``.
        windows : int
            Number of windows to be sampled.

        Returns
        -------
        schedule : list of floats
            A valid lambda schedule.
        """
        if schedule is None:
            return np.linspace(0., 1., windows)

        # Check end states
        if schedule[0] != 0 or schedule[-1] != 1:
            errmsg = ("end and start lambda windows must be lambda 0 and 1 "
                      "respectively")
            raise ValueError(errmsg)

        # Check monotonically increasing
        difference = np.diff(schedule)

        if not all(i >= 0. for i in difference):
            errmsg = "The lambda schedule is not monotonic"
            raise ValueError(errmsg)

        return schedule

    def _validate_functions(self, n=10):
        """Ensures that all the lambda functions adhere to the rules:
            - must begin at 0.
            - must finish at 1.
            - must be monotonically increasing

        Parameters
        ----------
        n : int, default 10
            number of grid points used to check monotonicity
        """
        # the individual lambda functions that must be defined for
        required_functions = list(LambdaProtocol.default_functions.keys())

        for function in required_functions:
            if function not in self.functions:
                # IA switched from warn to error here
                errmsg = (f"function {function} is missing from "
                          "self.lambda_functions.")
                raise ValueError(errmsg)

            # Check that the function starts and ends at 0 and 1 respectively
            if self.functions[function](0) != 0:
                raise ValueError("lambda functions must start at 0")
            if self.functions[function](1) != 1:
                raise ValueError("lambda fucntions must end at 1")

            # now validatate that it's monotonic
            global_lambda = np.linspace(0., 1., n)
            sub_lambda = [self.functions[function](lam) for
                          lam in global_lambda]
            difference = np.diff(sub_lambda)

            if not all(i >= 0. for i in difference):
                wmsg = (f"The function {function} is not monotonic as "
                        "typically expected.")
                warnings.warn(wmsg)

    def _check_for_naked_charges(self):
        """
        Checks that there are no cases where atoms have charge but no sterics.

        This avoids issues with singularities and/or excessive forces near
        the end states (even when using softcore electrostatics).
        """
        global_lambda = self.lambda_schedule

        def check_overlap(ele, sterics, global_lambda, functions, endstate):
            for lam in global_lambda:
                ele_val = functions[ele](lam)
                ster_val = functions[sterics](lam)
                # if charge > 0 and sterics == 0 raise error
                if ele_val != endstate and ster_val == endstate:
                    errmsg = ("There are states along this lambda schedule "
                              "where there are atoms with charges but no LJ "
                              f"interactions: {lam} {ele_val} {ster_val}")
                    raise ValueError(errmsg)

        # checking unique new terms first
        ele = 'lambda_electrostatics_insert'
        sterics = 'lambda_sterics_insert'
        check_overlap(ele, sterics, global_lambda, self.functions, endstate=0)

        # checking unique old terms now
        ele = 'lambda_electrostatics_delete'
        sterics = 'lambda_sterics_delete'
        check_overlap(ele, sterics, global_lambda, self.functions, endstate=1)

    def get_functions(self):
        return self.functions

    def plot_functions(self, lambda_schedule=None):
        """
        Plot the function for ease of visualisation.

        Parameters
        ----------
        shedule : np.ndarray
            The lambda schedule to plot the function along. If ``None`` plot
            the one stored within this class. Default ``None``.
        """
        import matplotlib.pyplot as plt

        fig = plt.figure(figsize=(10, 5))

        global_lambda = lambda_schedule if lambda_schedule else self.lambda_schedule

        for f in self.functions:
            plt.plot(global_lambda,
                     [self.functions[f](lam) for lam in global_lambda],
                     alpha=0.5, label=f)

        plt.xlabel('global lambda')
        plt.ylabel('sub-lambda')
        plt.legend()
        plt.show()

class RelativeAlchemicalState(AlchemicalState):
    """
    Relative AlchemicalState to handle all lambda parameters required for
    relative perturbations
    lambda = 1 refers to ON, i.e. fully interacting while
    lambda = 0 refers to OFF, i.e. non-interacting with the system
    all lambda functions will follow from 0 -> 1 following the master lambda
    lambda*core parameters perturb linearly
    lambda_sterics_insert and lambda_electrostatics_delete perturb in the
    first half of the protocol 0 -> 0.5
    lambda_sterics_delete and lambda_electrostatics_insert perturb in the
    second half of the protocol 0.5 -> 1

    Attributes
    ----------
    lambda_sterics_core
    lambda_electrostatics_core
    lambda_sterics_insert
    lambda_sterics_delete
    lambda_electrostatics_insert
    lambda_electrostatics_delete
    """

    class _LambdaParameter(AlchemicalState._LambdaParameter):
        pass

    lambda_sterics_core = _LambdaParameter('lambda_sterics_core')
    lambda_electrostatics_core = _LambdaParameter('lambda_electrostatics_core')
    lambda_sterics_insert = _LambdaParameter('lambda_sterics_insert')
    lambda_sterics_delete = _LambdaParameter('lambda_sterics_delete')
    lambda_electrostatics_insert = _LambdaParameter(
                                       'lambda_electrostatics_insert')
    lambda_electrostatics_delete = _LambdaParameter(
                                      'lambda_electrostatics_delete')

    def set_alchemical_parameters(self, global_lambda,
                                  lambda_protocol=LambdaProtocol()):
        """Set each lambda value according to the lambda_functions protocol.
        The undefined parameters (i.e. those being set to None) remain
        undefined.
        Parameters
        ----------
        lambda_value : float
            The new value for all defined parameters.
        """
        self.global_lambda = global_lambda
        for parameter_name in lambda_protocol.functions:
            lambda_value = lambda_protocol.functions[parameter_name](global_lambda)
            setattr(self, parameter_name, lambda_value)


#!/usr/bin/env python
# Copyright 2014-2019 The PySCF Developers. All Rights Reserved.
#
# 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.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

import os, sys
import imp
import numpy
import pyscf.ao2mo

try:
    from pyscf.dmrgscf import settings
    PyCheMPS2 = None
except ImportError:
    try:
        import PyCheMPS2
    except ImportError:
        msg = ('settings.py not found.  Please create %s\n' %
               os.path.join(os.path.dirname(__file__), 'settings.py'))
        sys.stderr.write(msg)

# point group ID defined in CheMPS2, see
# http://sebwouters.github.io/CheMPS2/classCheMPS2_1_1Irreps.html
GROUPNAME_ID = {
    'C1' : 0,
    'Ci' : 1,
    'C2' : 2,
    'Cs' : 3,
    'D2' : 4,
    'C2v': 5,
    'C2h': 6,
    'D2h': 7,
}

class CheMPS2(object):
    def __init__(self, mol, **kwargs):
        self.mol = mol
        self.verbose = mol.verbose
        self.stdout = mol.stdout
        if self.mol.symmetry:
            self.groupname = mol.groupname
        else:
            self.groupname = None
        self.orbsym = []

# ref.
# https://github.com/SebWouters/CheMPS2/blob/master/psi4plugins/dmrgci.cc
        self.wfn_irrep = 0
        #self.spin_2s = 0  # spin = 2*s, 0 means singlet
        self.dmrg_states = [ 200 , 500 , 1000 , 1000 ]
        self.dmrg_noise = [ 1 , 1 , 1 , 0 ]
        self.dmrg_e_convergence = 1e-8
        self.dmrg_noise_factor = 0.03
        self.dmrg_maxiter_noise = 5
        self.dmrg_maxiter_silent = 100

        self._keys = set(self.__dict__.keys())

    def dump_flags(self, verbose=None):
        log = pyscf.lib.logger.new_logger(self, verbose)
        log.info('******** CheMPS2 flags ********')
        log.info('dmrg_states = %s', str(self.dmrg_states))
        log.info('dmrg_noise = %s', str(self.dmrg_noise))
        log.info('dmrg_e_convergence = %g', self.dmrg_e_convergence)
        log.info('dmrg_noise_factor = %g', self.dmrg_noise_factor)
        log.info('dmrg_maxiter_noise = %d', self.dmrg_maxiter_noise)
        log.info('dmrg_maxiter_silent = %d', self.dmrg_maxiter_silent)

    def kernel(self, h1e, eri, norb, nelec, ci0=None, ecore=0, **kwargs):
        global PyCheMPS2
        if PyCheMPS2 is None:
            PyCheMPS2 = imp.load_dynamic('PyCheMPS2', settings.PYCHEMPS2BIN)

        Initializer = PyCheMPS2.PyInitialize()
        Initializer.Init()

        if self.groupname is not None:
            groupNumber = GROUPNAME_ID[self.groupname]
        else:
            groupNumber = 0
            self.orbsym = numpy.zeros(norb, numpy.int32)
        Ham = PyCheMPS2.PyHamiltonian(norb, groupNumber,
                                      numpy.asarray(self.orbsym, dtype=numpy.int32))
        eri = pyscf.ao2mo.restore(1, eri, norb)
        for i in range(norb):
            for j in range(norb):
                totsym = self.orbsym[i] ^ self.orbsym[j]
                if 0 == totsym:
                    Ham.setTmat(i, j, h1e[i,j])
                for k in range(norb):
                    for l in range(norb):
                        totsym = self.orbsym[i] \
                               ^ self.orbsym[j] \
                               ^ self.orbsym[k] \
                               ^ self.orbsym[l]
                        if 0 == totsym:
                            Ham.setVmat(i, k, j, l, eri[i,j,k,l])
        Ham.setEconst(0)

        if isinstance(nelec, (int, numpy.integer)):
            spin2 = 0
        else:
            spin2 = (nelec[0]-nelec[1])
            nelec = sum(nelec)

        Prob = PyCheMPS2.PyProblem(Ham, spin2, nelec, self.wfn_irrep)
        Prob.SetupReorderD2h()

        OptScheme = PyCheMPS2.PyConvergenceScheme(len(self.dmrg_states))
        for cnt,m in enumerate(self.dmrg_states):
            if self.dmrg_noise[cnt]:
                OptScheme.setInstruction(cnt, m, self.dmrg_e_convergence,
                                         self.dmrg_maxiter_noise,
                                         self.dmrg_noise_factor)
            else:
                OptScheme.setInstruction(cnt, m, self.dmrg_e_convergence,
                                         self.dmrg_maxiter_silent, 0.0)

        with pyscf.lib.capture_stdout() as stdout:
            theDMRG = PyCheMPS2.PyDMRG(Prob, OptScheme)
            Energy = theDMRG.Solve() + ecore
            theDMRG.calc2DMandCorrelations()
            pyscf.lib.logger.debug1(self.mol, stdout.read())

        rdm2 = numpy.empty((norb,)*4)
        for i in range(norb):
            for j in range(norb):
                for k in range(norb):
                    for l in range(norb):
                        rdm2[i,j,k,l] = theDMRG.get2DMA(i, j, k, l)

        with pyscf.lib.capture_stdout() as stdout:
            theDMRG.deleteStoredOperators()
            pyscf.lib.logger.debug1(self.mol, stdout.read())

# The order of deallocation matters!
        del(theDMRG)
        del(OptScheme)
        del(Prob)
        del(Ham)
        del(Initializer)

        fakewfn_by_rdm2 = rdm2
        return Energy, fakewfn_by_rdm2

    def make_rdm12(self, fakewfn_by_rdm2, ncas, nelec, **kwargs):
        if not isinstance(nelec, (int, numpy.integer)):
            nelec = sum(nelec)
# CheMPS2 uses physics notation
        rdm2 = fakewfn_by_rdm2.transpose(0,2,1,3)
        rdm1 = numpy.einsum('ijkk->ij', rdm2) / (nelec-1)
        return rdm1, rdm2

    def make_rdm1(self, fcivec, norb, nelec, link_index=None, **kwargs):
        return self.make_rdm12(fcivec, norb, nelec, **kwargs)[0]

if __name__ == '__main__':
    from pyscf import gto
    from pyscf import scf
    from pyscf import mcscf

    b = 1.4
    mol = gto.Mole()
    mol.build(
        verbose = 5,
        output = 'out-chemps2',
        atom = [['H', (0.,0.,i)] for i in range(8)],
        basis = {'H': 'sto-3g'},
        symmetry = True,
        symmetry_subgroup = 'D2h',
    )
    m = scf.RHF(mol)
    m.scf()

    mc = mcscf.CASSCF(m, 4, 4)
    mc.fcisolver = CheMPS2(mol)
    mc.fcisolver.dmrg_e_convergence = 1e-8
    emc_1 = mc.mc2step()[0]

    mc = mcscf.CASCI(m, 4, 4)
    mc.fcisolver = CheMPS2(mol)
    emc_0 = mc.casci()[0]

    b = 1.4
    mol = gto.Mole()
    mol.build(
        verbose = 5,
        output = 'out-casscf',
        atom = [['H', (0.,0.,i)] for i in range(8)],
        basis = {'H': 'sto-3g'},
        symmetry = True,
    )
    m = scf.RHF(mol)
    m.scf()

    mc = mcscf.CASSCF(m, 4, 4)
    emc_1ref = mc.mc2step()[0]

    mc = mcscf.CASCI(m, 4, 4)
    emc_0ref = mc.casci()[0]

    print('CheMPS2-CI  = %.15g CASCI  = %.15g' % (emc_0, emc_0ref))
    print('CheMPS2-SCF = %.15g CASSCF = %.15g' % (emc_1, emc_1ref))


"""
IMPORTANT : I have not written any exception handling here since this is just a prototype module.
            But exceptions can be logged into database with "traceback.format_exec()"
"""

import os
import numpy as np
from PIL import Image
import torch
import torch.nn as nn
import torchvision.models as models
from torchvision import transforms
from vocabulary import Vocabulary

# This python script will be called by Node.js as child process into project root directory, hence file paths will be from one level up.

ENCODER_CNN_CHECKPOINT = "python_models/saved_models/encoderEpoch_2.pth"
DECODER_LSTM_RNN_CHECKPOINT = "python_models/saved_models/decoderEpoch_2.pth"
VOCAB_FILE = "python_models/saved_models/vocab.pkl"

# Globally Set device to CPU since GPU won't be necessary for making just one prediction at a time
device = "cpu"

#=========================================================================
# Encoder - Decoder Model Class to be used for Generating Caption
#=========================================================================

class EncoderCNN(nn.Module):
    def __init__(self, embed_size):
        super(EncoderCNN, self).__init__()
        resnet = models.resnet50(pretrained=True)
        for param in resnet.parameters():
            param.requires_grad_(False)

        modules = list(resnet.children())[:-1]
        self.resnet = nn.Sequential(*modules)
        self.embed = nn.Linear(resnet.fc.in_features, embed_size)

    def forward(self, images):
        features = self.resnet(images)
        features = features.view(features.size(0), -1)
        features = self.embed(features)
        return features

class DecoderRNN(nn.Module):
    def __init__(self, embed_size, hidden_size, vocab_size, num_layers=1):
        super(DecoderRNN, self).__init__()
        self.embed_size = embed_size
        self.hidden_size = hidden_size
        self.vocab_size = vocab_size
        self.num_layers = num_layers

        self.embed = nn.Embedding(num_embeddings = self.vocab_size,
                                  embedding_dim = self.embed_size)
        self.lstm = nn.LSTM(input_size = self.embed_size,
                           hidden_size = self.hidden_size,
                           num_layers = self.num_layers,
                           batch_first = True)
        self.fc1 = nn.Linear(in_features = self.hidden_size,
                            out_features = self.vocab_size)

    def forward(self, features, captions):
        captions = captions[:,:-1]
        embeddings = self.embed(captions)
        embeddings = torch.cat((features.unsqueeze(1), embeddings),dim=1)
        batch_size = features.shape[0]

        self.hidden = (torch.zeros((self.num_layers, batch_size, self.hidden_size), device=device),torch.zeros((self.num_layers, batch_size, self.hidden_size), device=device))

        lstm_out, self.hidden = self.lstm(embeddings, self.hidden)
        output = self.fc1(lstm_out)
        return output

    def sample(self, inputs, states=None, max_len=20):
        " accepts pre-processed image tensor (inputs) and returns predicted sentence (list of tensor ids of length max_len) "

        output = []
        batch_size = inputs.shape[0]

        self.hidden = (torch.zeros((self.num_layers, batch_size, self.hidden_size), device=device),torch.zeros((self.num_layers, batch_size, self.hidden_size), device=device))

        while True:
            lstm_out, self.hidden = self.lstm(inputs, self.hidden)
            linear_out = self.fc1(lstm_out.squeeze(1))
            top_score = linear_out.max(1)[1]
            output.append(top_score.cpu().numpy()[0].item())
            if top_score == 1:
                break
            inputs = self.embed(top_score)
            inputs = inputs.unsqueeze(1)
        return output

def clean_sentence(output,vocab):
    cleaned = []
    for index in output:
        cleaned.append(vocab.idx2word[index])
    cleaned = cleaned[1:-1]
    sentence = ' '.join(cleaned)
    return sentence

def generateCaptions(image_path):
    vocab = Vocabulary(vocab_file=VOCAB_FILE,vocab_from_file=True,vocab_threshold=None)

    # Model parameters
    embed_size = 512
    hidden_size = 512
    vocab_size = len(vocab)

    # Initialize the encoder and decoder, and set each to inference mode.
    encoder = EncoderCNN(embed_size)
    decoder = DecoderRNN(embed_size, hidden_size, vocab_size, num_layers=2)

    # Load the trained weights.
    encoder.load_state_dict(torch.load(ENCODER_CNN_CHECKPOINT))
    decoder.load_state_dict(torch.load(DECODER_LSTM_RNN_CHECKPOINT))

    # Put model to evaluation mode
    encoder.eval()
    decoder.eval()

    # Move models to GPU if CUDA is available.
    encoder.to(device)
    decoder.to(device)

    transform_predict = transforms.Compose([
    transforms.Resize(256),                          # smaller edge of image resized to 256
    transforms.RandomCrop(224),                      # get 224x224 crop from random location
    transforms.RandomHorizontalFlip(),               # horizontally flip image with probability=0.5
    transforms.ToTensor(),                           # convert the PIL Image to a tensor
    transforms.Normalize((0.485, 0.456, 0.406),      # normalize image for pre-trained model
                         (0.229, 0.224, 0.225))])

    image = Image.open(image_path).convert('RGB')    # Convert to RGB in case additonal channel is there in image
    original_image = np.copy(image)
    image = torch.unsqueeze(transform_predict(image), 0).to(device)
    features = encoder(image).unsqueeze(1)
    output = decoder.sample(features)
    sentence = clean_sentence(output, vocab)
    return {"sentence":sentence}


"""
=======
fourier
=======

This module define the Fourier class and other functions related to
the fourier transformation.
"""
import numpy as np
import pandas as pd
from javelin.grid import Grid
from javelin.utils import get_unitcell, get_positions, get_atomic_numbers
from javelin.fourier_cython import calculate_cython, approx_calculate_cython

class Fourier:
    """The Fourier class contains everything required to calculate the
    diffuse scattering. The only required thing to be set is
    :obj:`javelin.fourier.Fourier.structure`. There are defaults for
    all other options including **grid**, **radiation**, **average**
    structure subtraction and **lots** options.

    :examples:

    >>> from javelin.structure import Structure
    >>> fourier = Fourier()
    >>> print(fourier)
    Radiation         : neutron
    Fourier volume    : complete crystal
    Aver. subtraction : False
    <BLANKLINE>
    Reciprocal layer  :
    lower left  corner :     [ 0.  0.  0.]
    lower right corner :     [ 1.  0.  0.]
    upper left  corner :     [ 0.  1.  0.]
    top   left  corner :     [ 0.  0.  1.]
    <BLANKLINE>
    hor. increment     :     [ 0.01  0.    0.  ]
    vert. increment    :     [ 0.    0.01  0.  ]
    top   increment    :     [ 0.  0.  1.]
    <BLANKLINE>
    # of points        :     101 x 101 x 1
    >>> results = fourier.calc(Structure())
    >>> print(results) # doctest: +SKIP
    <xarray.DataArray 'Intensity' ([ 1.  0.  0.]: 101, [ 0.  1.  0.]: 101, [ 0.  0.  1.]: 1)>
    array([[[ 0.],
            [ 0.],
            ...,
            [ 0.],
            [ 0.]],
    <BLANKLINE>
           [[ 0.],
            [ 0.],
            ...,
            [ 0.],
            [ 0.]],
    <BLANKLINE>
           ...,
           [[ 0.],
            [ 0.],
            ...,
            [ 0.],
            [ 0.]],
    <BLANKLINE>
           [[ 0.],
            [ 0.],
            ...,
            [ 0.],
            [ 0.]]])
    Coordinates:
      * [ 1.  0.  0.]  ([ 1.  0.  0.]) float64 0.0 0.01 0.02 0.03 0.04 0.05 0.06 ...
      * [ 0.  1.  0.]  ([ 0.  1.  0.]) float64 0.0 0.01 0.02 0.03 0.04 0.05 0.06 ...
      * [ 0.  0.  1.]  ([ 0.  0.  1.]) float64 0.0
    Attributes:
        units:    r.l.u

    """
    def __init__(self):
        self._radiation = 'neutron'
        self._lots = None
        self._number_of_lots = None
        self._average = False
        self._magnetic = False
        self._cython = True
        self._approximate = True
        self._fast = True

        #: The **grid** attribute defines the reciprocal volume from
        #: which the scattering will be calculated. Must of type
        #: :class:`javelin.grid.Grid` And check
        #: :class:`javelin.grid.Grid` for details on how to change the
        #: grid.
        self.grid = Grid()

    def __str__(self):
        return """Radiation         : {}
Fourier volume    : {}
Aver. subtraction : {}

Reciprocal layer  :
{}""".format(self.radiation,
             "complete crystal" if self.lots is None else "{} lots of {} x {} x {} unit cells"
             .format(self.number_of_lots, *self.lots),
             self.average,
             self.grid)

    @property
    def radiation(self):
        """The radiation used

        :getter: Returns the radiation selected
        :setter: Sets the radiation
        :type: str ('xray' or 'neutron')
        """
        return self._radiation

    @radiation.setter
    def radiation(self, rad):
        if rad not in ('neutron', 'xray'):
            raise ValueError("radiation must be one of 'neutron' or 'xray'")
        self._radiation = rad

    @property
    def lots(self):
        """The size of lots

        :getter: Returns the lots size
        :setter: Sets the lots size
        :type: list of 3 integers or None
        """
        return self._lots

    @lots.setter
    def lots(self, lots):
        if lots is None:
            self._lots = None
        else:
            lots = np.asarray(lots)
            if len(lots) == 3:
                self._lots = lots
            else:
                raise ValueError("Must provied 3 values for lots")

    @property
    def number_of_lots(self):
        """The number of lots to use

        :getter: Returns the number of lots
        :setter: Sets the number of lots
        :type: int
        """
        return self._number_of_lots

    @number_of_lots.setter
    def number_of_lots(self, value):
        self._number_of_lots = value

    @property
    def average(self):
        """This sets the options of calculating average structure and
        subtracted it from the simulated scattering

        :getter: Returns bool of average structure subtraction option
        :setter: Sets whether average structure is subtracted
        :type: bool
        """
        return self._average

    @average.setter
    def average(self, value):
        if isinstance(value, bool):
            self._average = value
        else:
            raise TypeError("Expected a bool, True or False")

    @property
    def magnetic(self):
        """This sets the options of calculating the magnetic scattering
        instead of nuclear. This assume neutrons are being used.

        :getter: Returns bool of magnetic scattering option
        :setter: Sets whether magnetic sacttering is calculated
        :type: bool

        """
        return self._magnetic

    @magnetic.setter
    def magnetic(self, value):
        if isinstance(value, bool):
            self._magnetic = value
        else:
            raise TypeError("Expected a bool, True or False")

    @property
    def approximate(self):
        """This sets the options of calculating the approximate scattering
        instead of exact. This is much quicker and is likely good enough for
        most cases.

        :getter: Returns bool of approximate scattering option
        :setter: Sets whether approximate sacttering is calculated
        :type: bool

        """
        return self._approximate

    @approximate.setter
    def approximate(self, value):
        if isinstance(value, bool):
            self._approximate = value
        else:
            raise TypeError("Expected a bool, True or False")

    def __get_q(self, unitcell):
        qx, qy, qz = self.grid.get_q_meshgrid()
        q = np.linalg.norm(np.array([qx.ravel(),
                                     qy.ravel(),
                                     qz.ravel()]).T @ unitcell.B, axis=1)
        q.shape = qx.shape
        return q*2*np.pi

    def calc(self, structure):
        """Calculates the fourier transform

        :param structure: The structure from which fourier transform
        is calculated. The calculation work with any of the following
        types of structures :class:`javelin.structure.Structure`,
        :class:`ase.Atoms` or
        :class:`diffpy.Structure.structure.Structure` but if you are
        using average structure subtraction or the lots option it
        needs to be :class:`javelin.structure.Structure` type.

        :return: DataArray containing calculated diffuse scattering
        :rtype: :class:`xarray.DataArray`

        """

        if structure is None:
            raise ValueError("You have not set a structure for this calculation")

        if self.average:
            aver = self._calculate_average(structure)

        unitcell = get_unitcell(structure)

        if self.lots is None:
            atomic_numbers = get_atomic_numbers(structure)
            positions = get_positions(structure)
            if self.magnetic:
                magmons = structure.get_magnetic_moments()
                return create_xarray_dataarray(self._calculate_magnetic(atomic_numbers,
                                                                        positions,
                                                                        unitcell,
                                                                        magmons), self.grid)
            else:
                results = self._calculate(atomic_numbers, positions, unitcell)
                if self.average:
                    results -= aver

                return create_xarray_dataarray(np.real(results*np.conj(results)), self.grid)

        else:  # needs to be Javelin structure, lots by unit cell
            total = np.zeros(self.grid.bins)
            levels = structure.atoms.index.levels
            for lot in range(self.number_of_lots):
                print(lot+1, 'out of', self.number_of_lots)
                starti = np.random.randint(len(levels[0]))
                startj = np.random.randint(len(levels[1]))
                startk = np.random.randint(len(levels[2]))
                ri = np.roll(levels[0], -starti)[:self.lots[0]]
                rj = np.roll(levels[1], -startj)[:self.lots[1]]
                rk = np.roll(levels[2], -startk)[:self.lots[2]]
                atoms = structure.atoms.loc[ri, rj, rk, :]
                atomic_numbers = atoms.Z.values
                positions = (atoms[['x', 'y', 'z']].values +
                             np.asarray([np.mod(atoms.index.get_level_values(0).values-starti,
                                                len(levels[0])),
                                         np.mod(atoms.index.get_level_values(1).values-startj,
                                                len(levels[1])),
                                         np.mod(atoms.index.get_level_values(2).values-startk,
                                                len(levels[2]))]).T)
                if self.magnetic:
                    magmons = structure.magmons.loc[ri, rj, rk, :].values
                    total += self._calculate_magnetic(atomic_numbers, positions, unitcell, magmons)
                else:
                    results = self._calculate(atomic_numbers, positions, unitcell)
                    if self.average:
                        results -= aver
                    total += np.real(results*np.conj(results))

            scale = (structure.atoms.index.droplevel(3).drop_duplicates().size /
                     (self.number_of_lots*self.lots.prod()))

            return create_xarray_dataarray(total*scale, self.grid)

    def calc_average(self, structure):
        """Calculates the scattering from the avarage structure

        :param structure: The structure from which fourier transform
        is calculated. The calculation work with any of the following
        types of structures :class:`javelin.structure.Structure`,
        :class:`ase.Atoms` or
        :class:`diffpy.Structure.structure.Structure` but if you are
        using average structure subtraction or the lots option it
        needs to be :class:`javelin.structure.Structure` type.

        :return: DataArray containing calculated average scattering
        :rtype: :class:`xarray.DataArray`
        """

        if structure is None:
            raise ValueError("You have not set a structure for this calculation")

        aver = self._calculate_average(structure)
        return create_xarray_dataarray(np.real(aver*np.conj(aver)), self.grid)

    def _calculate_average(self, structure):
        aver = self._calculate(get_atomic_numbers(structure),
                               structure.xyz, get_unitcell(structure))

        aver /= structure.atoms.index.droplevel(3).drop_duplicates().size

        # compute the interference function of the lot shape

        if self.lots is None:
            index = structure.atoms.index.droplevel(3).drop_duplicates()
        else:
            index = structure.atoms.loc[pd.RangeIndex(self.lots[0]),
                                        pd.RangeIndex(self.lots[1]),
                                        pd.RangeIndex(self.lots[2]),
                                        :].index.droplevel(3).drop_duplicates()

        aver *= self._calculate(np.zeros(len(index), dtype=np.int),
                                np.asarray([index.get_level_values(0).astype('double').values,
                                            index.get_level_values(1).astype('double').values,
                                            index.get_level_values(2).astype('double').values]).T,
                                use_ff=False)

        return aver

    def _calculate(self, atomic_numbers, positions, unitcell=None, use_ff=True):
        if self._fast and not self._cython:
            qx, qy, qz = self.grid.get_squashed_q_meshgrid()
        else:
            qx, qy, qz = self.grid.get_q_meshgrid()
        qx *= (2*np.pi)
        qy *= (2*np.pi)
        qz *= (2*np.pi)

        # Get unique list of atomic numbers
        unique_atomic_numbers = np.unique(atomic_numbers)

        results = np.zeros(self.grid.bins, dtype=np.complex)
        # Loop of atom types
        for atomic_number in unique_atomic_numbers:
            try:
                ff = get_ff(atomic_number, self.radiation, self.__get_q(unitcell)) if use_ff else 1
            except KeyError as e:
                print("Skipping fourier calculation for atom " + str(e) +
                      ", unable to get scattering factors.")
                continue

            atom_positions = positions[np.where(atomic_numbers == atomic_number)]
            temp_array = np.zeros(self.grid.bins, dtype=np.complex)
            print("Working on atom number", atomic_number, "Total atoms:", len(atom_positions))

            # Loop over atom positions of type atomic_number
            if self._cython:
                if self.approximate:
                    cex = np.exp(np.linspace(0, 2j*np.pi*(1-2**-16), 2**16))
                    approx_calculate_cython(self.grid.ll, self.grid.v1_delta,
                                            self.grid.v2_delta, self.grid.v3_delta,
                                            atom_positions, temp_array.real,
                                            temp_array.imag, cex.real, cex.imag)
                else:
                    calculate_cython(qx, qy, qz, atom_positions, temp_array.real, temp_array.imag)
            else:
                if self._fast:
                    for atom in atom_positions:
                        dotx = np.exp(qx*atom[0]*1j)
                        doty = np.exp(qy*atom[1]*1j)
                        dotz = np.exp(qz*atom[2]*1j)
                        temp_array += dotx * doty * dotz
                else:
                    for atom in atom_positions:
                        dot = qx*atom[0] + qy*atom[1] + qz*atom[2]
                        temp_array += np.exp(dot*1j)

            results += temp_array * ff  # scale by form factor

        return results

    def _calculate_magnetic(self, atomic_numbers, positions, unitcell, magmons):
        if self._fast:
            qx, qy, qz = self.grid.get_squashed_q_meshgrid()
        else:
            qx, qy, qz = self.grid.get_q_meshgrid()
        qx *= (2*np.pi)
        qy *= (2*np.pi)
        qz *= (2*np.pi)
        q2 = self.__get_q(unitcell)**2

        # Get unique list of atomic numbers
        unique_atomic_numbers = np.unique(atomic_numbers)

        # Loop of atom types
        spinx = np.zeros(self.grid.bins, dtype=np.complex)
        spiny = np.zeros(self.grid.bins, dtype=np.complex)
        spinz = np.zeros(self.grid.bins, dtype=np.complex)
        for atomic_number in unique_atomic_numbers:
            try:
                ff = get_mag_ff(atomic_number, self.__get_q(unitcell), ion=3)
            except (AttributeError, KeyError) as e:
                print("Skipping fourier calculation for atom " + str(e) +
                      ", unable to get magnetic scattering factors.")
                continue

            atom_positions = positions[np.where(atomic_numbers == atomic_number)]
            temp_spinx = np.zeros(self.grid.bins, dtype=np.complex)
            temp_spiny = np.zeros(self.grid.bins, dtype=np.complex)
            temp_spinz = np.zeros(self.grid.bins, dtype=np.complex)
            print("Working on atom number", atomic_number, "Total atoms:", len(atom_positions))
            # Loop over atom positions of type atomic_number
            if self._fast:
                for atom, spin in zip(atom_positions, magmons):
                    dotx = np.exp(qx*atom[0]*1j)
                    doty = np.exp(qy*atom[1]*1j)
                    dotz = np.exp(qz*atom[2]*1j)
                    exp_temp = dotx * doty * dotz
                    temp_spinx += exp_temp*spin[0]
                    temp_spiny += exp_temp*spin[1]
                    temp_spinz += exp_temp*spin[2]
            else:
                for atom, spin in zip(atom_positions, magmons):
                    dot = qx*atom[0] + qy*atom[1] + qz*atom[2]
                    exp_temp = np.exp(dot*1j)
                    temp_spinx += exp_temp*spin[0]
                    temp_spiny += exp_temp*spin[1]
                    temp_spinz += exp_temp*spin[2]
            spinx += temp_spinx * ff
            spiny += temp_spiny * ff
            spinz += temp_spinz * ff
        # Calculate vector rejection of spin onto q
        # M - M.Q/|Q|^2 Q
        scale = (spinx*qx + spiny*qy + spinz*qz)/q2
        spinx = spinx - scale * qx
        spiny = spiny - scale * qy
        spinz = spinz - scale * qz
        return np.real(spinx*np.conj(spinx) + spiny*np.conj(spiny) + spinz*np.conj(spinz))

def create_xarray_dataarray(values, grid):
    """Create a xarry DataArray from the input numpy array and grid
    object.

    :param values: Input array containing the scattering intensities
    :type values: :class:`numpy.ndarray`
    :param numbers: Grid object describing the array properties
    :type numbers: :class:`javelin.grid.Grid`
    :return: DataArray produced from the values and grid object
    :rtype: :class:`xarray.DataArray`
    """
    import xarray as xr
    return xr.DataArray(data=values,
                        name="Intensity",
                        dims=(grid.get_axes_names()),
                        coords=(grid.r1, grid.r2, grid.r3),
                        attrs=(("units", grid.units),))

def get_ff(atomic_number, radiation, q=None):
    """Returns the form factor for a given atomic number, radiation and q
    values

    :param atomic_number: atomic number
    :type atomic_number: int
    :param radiation: type of radiation ('xray' or 'neutron')
    :type radiation: str
    :param q: value or values of q for which to get form factors
    :type q: float, list, :class:`numpy.ndarray`
    :return: form factors for given q
    :rtype: float, :class:`numpy.ndarray`

    :Examples:

    >>> get_ff(8, 'neutron')
    5.805

    >>> get_ff(8, 'xray', q=2.0)
    6.31826029176493

    >>> get_ff(8, 'xray', q=[0.0, 3.5, 7.0])
    array([ 7.999706  ,  4.38417867,  2.08928068])
    """
    import periodictable

    if atomic_number < 1:
        raise KeyError(atomic_number)

    if radiation == 'neutron':
        return periodictable.elements[atomic_number].neutron.b_c
    elif radiation == 'xray':
        return periodictable.elements[atomic_number].xray.f0(q)
    else:
        raise ValueError("Unknown radition: " + radiation)

def get_mag_ff(atomic_number, q, ion=0, j=0):
    """Returns the j0 magnetic form factor for a given atomic number,
    radiation and q values

    :param atomic_number: atomic number
    :type atomic_number: int
    :param q: value or values of q for which to get form factors
    :type q: float, list, :class:`numpy.ndarray`
    :param ion: charge of selected atom
    :type ion: int
    :param j: order of spherical Bessel function (0, 2, 4 or 6)
    :type j: int
    :return: magnetic form factor for given q
    :rtype: float, :class:`numpy.ndarray`

    :Examples:

    >>> get_mag_ff(8, q=2, ion=1)
    0.58510426376585045

    >>> get_mag_ff(26, q=[0.0, 3.5, 7.0], ion=2)
    array([ 1.        ,  0.49729671,  0.09979243])

    >>> get_mag_ff(26, q=[0.0, 3.5, 7.0], ion=4)
    array([ 0.9997    ,  0.58273549,  0.13948496])

    >>> get_mag_ff(26, q=[0.0, 3.5, 7.0], ion=4, j=4)
    array([ 0.       ,  0.0149604,  0.0759222])
    """
    import periodictable
    return getattr(periodictable.elements[atomic_number].magnetic_ff[ion], 'j'+str(j)+'_Q')(q)


#!/usr/bin/env python
# -*- coding: utf-8 -*-

##############
# GaudiMM: Genetic Algorithms with Unrestricted
# Descriptors for Intuitive Molecular Modeling
#
# https://github.com/insilichem/gaudi
#
# Copyright 2017 Jaime Rodriguez-Guerra, Jean-Didier Marechal
#
# 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.
##############

"""
This objective performs rough estimations
of good orientations of ligating residues in a protein to
coordinate a given metal or small molecule. The geometry is approximated
by computing average distances from ligating atoms the metal centre (``self.probe``)
as well as the angles formed by the probe, the ligating atom and its immediate neighbor.
Good planarity is assured by a dihedral check.

"""

# Python
from __future__ import print_function, division
import math
import logging
import numpy as np
# Chimera
import chimera
# GAUDI
from gaudi.objectives import ObjectiveProvider
from gaudi.exceptions import AtomsNotFound, ResiduesNotFound, MoleculesNotFound
from gaudi import parse
from gaudi._cpdrift import coherent_point_drift

logger = logging.getLogger(__name__)

GEOMETRIES = {
    'cube': np.array([(1.0, 1.0, 1.0), (1.0, 1.0, -1.0), (1.0, -1.0, 1.0), (1.0, -1.0, -1.0), (-1.0, 1.0, 1.0), (-1.0, 1.0, -1.0), (-1.0, -1.0, 1.0), (-1.0, -1.0, -1.0), (0.0, 0.0, 0.0)]),
    'hexagonal bipyramid': np.array([(0.0, 1.0, 0.0), (0.0, -1.0, 0.0), (0.0, 0.0, 1.0), (0.0, 0.0, -1.0), (0.8660, 0.0, 0.5), (0.8660, 0.0, -0.5), (-0.8660, 0.0, 0.5), (-0.8660, 0.0, -0.5), (0.0, 0.0, 0.0)]),
    'linear': np.array([(0.0, 1.0, 0.0), (0.0, -1.0, 0.0), (0.0, 0.0, 0.0)]),
    'octahedron': np.array([(1.0, 0.0, 0.0), (-1.0, 0.0, 0.0), (0.0, 1.0, 0.0), (0.0, -1.0, 0.0), (0.0, 0.0, 1.0), (0.0, 0.0, -1.0), (0.0, 0.0, 0.0)]),
    'pentagonal bipyramid': np.array([(0.0, 1.0, 0.0), (0.0, -1.0, 0.0), (0.0, 0.0, 1.0), (0.9511, 0.0, 0.30901), (0.5878, 0.0, -0.8090), (-0.5878, 0.0, -0.8090), (-0.9511, 0.0, 0.3090), (0.0, 0.0, 0.0)]),
    'square planar': np.array([(1.0, 0.0, 0.0), (-1.0, 0.0, 0.0), (0.0, 1.0, 0.0), (0.0, -1.0, 0.0), (0.0, 0.0, 0.0)]),
    'square pyramid': np.array([(0.0, 1.0, 0.0), (-0.6123, -0.5, 0.6123), (0.6123, -0.5, 0.6123), (-0.6123, -0.5, -0.6123), (0.6123, -0.5, -0.6123), (0.0, 0.0, 0.0)]),
    'tetrahedral': np.array([(1.0, 1.0, 1.0), (-1.0, -1.0, 1.0), (-1.0, 1.0, -1.0), (1.0, -1.0, -1.0), (0.0, 0.0, 0.0)]),
    'trigonal planar': np.array([(1.0, 0.0, 0.0), (-0.5, 0.0, 0.8660), (-0.5, 0.0,-0.8660), (0.0, 0.0, 0.0)]),
    'trigonal bipyramid': np.array([(0.0, 1.0, 0.0), (0.0, -1.0, 0.0), (1.0, 0.0, 0.0), (-0.5, 0.0, 0.8660), (-0.5, 0.0, -0.8660), (0.0, 0.0, 0.0)]),
    'trigonal prism': np.array([(-0.6547, 0.6547, 0.3779), (0.6547, 0.6547, 0.3779), (0.0, 0.6547, -0.7559), (-0.6547, -0.6547, 0.3779), (0.6547, -0.6547, 0.3779), (0.0, -0.6547, -0.7559), (0.0, 0.0, 0.0)])}

def enable(**kwargs):
    kwargs = Coordination.validate(kwargs)
    return Coordination(**kwargs)

class Coordination(ObjectiveProvider):

    """
    Coordination class

    Parameters
    ----------
    probe : tuple
        The atom that acts as the metal center, expressed as
        <molecule_name>/<atom serial>. This will be parsed later on.
    residues : list of str
        Residues that must coordinate to `probe`, expressed as
        <molecule_name>/<residue position>. Position can be `*`.
    radius : float, optional, default=3.0
        Distance from `probe` where ligating atoms must be found
    atom_types : list of str, optional
        Types of atoms that are considered ligands to `probe`
    atom_names : list of str, optional
        Names of atoms that are considered ligands to `probe`
    atom_elements : list of str, optional
        Elements of atoms that are considered ligands to `probe`
    distance : float, optional
        Perfect distance a ligand atom should be from target.
    geometry : str or list of 3-tuple floats, optional
        Which geometry should be fitted. Choose from `GEOMETRIES` dict or
        specify a set of vectors.
    enforce_all_residues: bool, optional
        Whether to force or not if all specified residues should coordinate.
    only_one_ligand_per_residue : bool, optional
        Enforce that only one ligand for each residue should coordinate.
    prevent_intruders : bool, optional
        Don't let non-ligand atoms to be closer to the target than the
        selected ligand atoms.
    center_of_mass_correction : bool, optional
        If True, calculate the distance between the metal center
        and the center of mass of the ligand atoms, and sum that
        to the final score.
    distance_correction : bool, optional
        If True, report the deviation of the experimental coordination
        bond length and the ideal one, as tabulated by ``chimera.Element``,
        and sum that to the final score.

    Returns
    -------
    float
        Sum of RMSD of vertices from ideal RMSD and average cosine of
        angle deviation from ideal orientation of ligand neighbors.
        A perfect match should report 0.0.
    """
    _validate = {
        parse.Required('probe'): parse.Named_spec("molecule", "atom"),
        parse.Required('residues'): [parse.Named_spec("molecule", "residue")],
        'radius': parse.Coerce(float),
        'atom_types': [basestring],
        'atom_names': [basestring],
        'atom_elements': [basestring],
        'distance': parse.All(parse.Coerce(float), parse.Range(min=0)),
        'min_atoms': parse.All(parse.Coerce(int), parse.Range(min=2)),
        'geometry': parse.Any(parse.In(GEOMETRIES.keys()), [parse.Coordinates]),
        'enforce_all_residues': parse.Coerce(bool),
        'only_one_ligand_per_residue': parse.Coerce(bool),
        'prevent_intruders': parse.Coerce(bool),
        'center_of_mass_correction': parse.Coerce(bool),
        'distance_correction': parse.Coerce(bool)
        }

    def __init__(self, probe=None, radius=3.0, atom_types=(), atom_elements=(),
                 atom_names=(), residues=(), geometry='tetrahedral', distance=0,
                 min_atoms=1, prevent_intruders=True, enforce_all_residues=False,
                 only_one_ligand_per_residue=False, center_of_mass_correction=False,
                 distance_correction=False, *args, **kwargs):
        ObjectiveProvider.__init__(self, **kwargs)
        self._probe = probe
        self._residues = residues
        self.radius = radius
        self.atom_types = atom_types
        self.atom_names = atom_names
        self.atom_elements = atom_elements
        self.distance = distance
        self.min_atoms = min_atoms
        self.only_one_ligand_per_residue = only_one_ligand_per_residue
        self.enforce_all_residues = enforce_all_residues
        self.prevent_intruders = prevent_intruders
        self.center_of_mass_correction = center_of_mass_correction
        self.distance_correction = distance_correction
        if isinstance(geometry, basestring):
            self.geometry = np.copy(GEOMETRIES[geometry])
        else:
            self.geometry = np.array(geometry)
        self.n_vertices = self.geometry.shape[0] - 1
        if self.n_vertices < self.min_atoms:
            self.min_ligands = self.n_vertices
            logger.warn('# Vertices in selected geometry < min_ligands! Overriding '
                        'min_ligands with {}'.format(self.n_vertices))
        if not sum(map(bool, [atom_types, atom_elements, atom_names])):
            raise ValueError('At least one of atom_types, atom_elements, atom_names '
                             'must be specified')

    def molecules(self, ind):
        return [m.compound.mol for m in ind._molecules.values()]

    def probe(self, ind):
        mol, serial = self._probe
        return ind.find_molecule(mol).find_atom(serial)

    def residues(self, ind):
        for mol, pos in self._residues:
            for residue in ind.find_molecule(mol).find_residues(pos):
                yield residue

    def evaluate(self, ind):
        """
        1. Get requested atoms sorted by distance
        2. If they meet the minimum quantity, return the rmsd for
           that geometry
        3. If that's not possible of they are not enough, return penalty
        """
        try:
            test_atoms = [a for d, a in self.coordination_sphere(ind)]
        except ResiduesNotFound:
            logger.warning("Not enough atoms or some residues missing")
            return -1000 * self.weight

        missing_atoms = self.min_atoms - len(test_atoms)
        if missing_atoms > 0:
            logger.warning("Could not find enough ligand atoms in probe environment. "
                           "{} missing".format(missing_atoms))
            return -100 * missing_atoms * self.weight

        geometry = np.copy(self.geometry)
        ligands = test_atoms[:self.n_vertices]
        metal = self.probe(ind)
        atom_points = [a.xformCoord() for a in ligands + [metal]]
        atom_points_array = np.array(atom_points)
        # rmsd
        try:
            _, _, rmsd = coherent_point_drift(atom_points_array, geometry, method='rigid',
                                              guess_steps=2, max_iterations=10)
        except Exception as e:
            logger.exception(e)  #
            logger.warning("Geometry not feasible in current conditions")
            return -1000 * self.weight
        # directionality
        directionality = sum(ideal_bond_deviation(metal, ligand, ligands) for ligand in ligands)

        com_deviation = 0
        if self.center_of_mass_correction:
            ligands_com = np.average(atom_points_array[:-1], axis=0)
            com_deviation = atom_points[-1].distance(chimera.Point(*ligands_com))

        dist_deviation = 0
        if self.distance_correction:
            for ligand in ligands:
                experimental_bond_length = ligand.xformCoord().distance(atom_points[-1])
                ideal_bond_length = chimera.Element.bondLength(ligand.element, metal.element)
                dist_deviation += abs(experimental_bond_length - ideal_bond_length)

        return rmsd + directionality + com_deviation + dist_deviation

    def coordination_sphere(self, ind):
        """
        1. Get atoms and residues found within ``self.radius``
        angstroms from ``self.probe``. Found residues MUST
        include ``self.residues``. Otherwise, apply penalty.

        2. Sort atoms by absolute difference of ``self.distance``
        and distance to ``self.probe``. That way, nearest atoms are computed first.
        If found atoms do not include some of the requested types, apply penalty.
        """
        # Helpers
        def abs_distance(a):
            return abs(self.distance - metal.xformCoord().distance(a.xformCoord()))
        if self.atom_types:
            def atom_is_valid(a): return a.idatmType in self.atom_types
        elif self.atom_elements:
            def atom_is_valid(a): return a.element.name in self.atom_elements
        elif self.atom_names:
            def atom_is_valid(a): return a.name in self.atom_names
        else:
            def atom_is_valid(a): return True

        self._update_zone(ind)
        metal = self.probe(ind)
        residues = list(self.residues(ind))
        atoms = [a for a in self.zone.atoms() if a is not metal]

        atoms_by_distance = []
        found_residues = set()
        distance_and_atoms = sorted((abs_distance(a), a) for a in atoms)
        for d, a in distance_and_atoms:
            if atom_is_valid(a) and a.residue in residues and d > 1.0:
                atoms_by_distance.append((d, a))
                found_residues.add(a.residue)
            elif self.prevent_intruders:
                break

        if self.enforce_all_residues and found_residues != residues:
            logger.warning("Some atoms found, but some residues are missing")
            raise ResiduesNotFound

        return atoms_by_distance

    def _update_zone(self, ind):
        """
        Clear existing selection and add atoms within `self.radius` from
        `self.probe`, as long as they belong to one of `self.molecules`
        """
        self.zone.clear()
        self.zone.add(self.probe(ind))
        self.zone.merge(chimera.selection.REPLACE,
                        chimera.specifier.zone(
                            self.zone, 'atom', None, self.radius, self.molecules(ind)))
        return self.zone

def ideal_bond_deviation(metal, ligand, other_ligands=()):

    """
    Assess if the current bond vector is well oriented with
    respect to the ideal bond vector.

    Parameters
    ----------
    metal : chimera.Atom
        The ion `ligands` are coordinating to
    ligand : chimera.Atom
        Potential ligand atoms to `metal`

    Returns
    -------
    float
        Absolute sine of the angle between the ideal vector and
        the ligand-metal one.
    """
    ligand_idatm = chimera.idatm.typeInfo.get(ligand.idatmType)
    ligand_geometry = ligand_idatm.geometry if ligand_idatm else 3
    ideal_positions = ideal_bonded_positions(ligand, metal.element, geometry=ligand_geometry)
    if not ideal_positions:
        logger.warning("Ligand %s reports no available bonding positions. Check your atom_types!", ligand)
        return 1.0
    # else we can go on
    # Conditions and booleans
    rotates = ligand_geometry == 4 and len(ligand.neighbors) == 1
    bidentate = False
    for bidentate_mate in other_ligands:
        if bidentate_mate is not ligand:
            shared_neighbors = set(ligand.neighbors) & set(bidentate_mate.neighbors)
            if shared_neighbors:
                bidentate = True
                break

    # Coordinates and positions
    ligand_coord = ligand.xformCoord()
    metal_coord = metal.xformCoord()
    if not ligand.neighbors:
        # If ligand has no neighbors, it's an isolated atom. This means it will always
        # be well oriented towards the metal. Then we can just return 0.0
        return 0.0
    # else we can go on
    neighbor = ligand.neighbors[0]
    neighbor_coord = neighbor.xformCoord()
    try:
        n_neighbor = next(a for a in neighbor.neighbors if a is not ligand)
        n_neighbor_coord = n_neighbor.xformCoord()
    except StopIteration:
        logger.warning('Ligand %s has no 2nd level neighbors. Dihedrals cannot be evaluated', ligand)
        rotates = True # This will force to ignore dihedral calculation ;)

    ideal_pos = ideal_positions[0]
    actual_angle = chimera.angle(neighbor_coord, ligand_coord, metal_coord)
    ideal_angles = [chimera.angle(neighbor_coord, ligand_coord, ideal_pos)]
    ideal_dihedrals = []
    if bidentate and not rotates:
        a = ligand_coord
        b = bidentate_mate.xformCoord()
        c = shared_neighbors.pop().xformCoord()
        bidentate_ideal_pos = c + 1.5 * ((a-c) + (b-c))
        bidentate_ideal_angle = chimera.angle(neighbor_coord, ligand_coord, bidentate_ideal_pos)
        bidentate_ideal_dihedral = chimera.dihedral(n_neighbor_coord, neighbor_coord,
                                                    ligand_coord, bidentate_ideal_pos)
        ideal_angles.append(bidentate_ideal_angle)
        ideal_dihedrals.append(bidentate_ideal_dihedral)

    angle_diff = min(delta - actual_angle for delta in ideal_angles)
    abs_sin_angle = abs(math.sin(math.radians(angle_diff)))
    if rotates:  # we don't care about the dihedral in this case
        return abs_sin_angle
    # else:
    actual_dihedral = chimera.dihedral(n_neighbor_coord, neighbor_coord, ligand_coord, metal_coord)
    ideal_dihedral = chimera.dihedral(n_neighbor_coord, neighbor_coord, ligand_coord, ideal_pos)
    ideal_dihedrals.append(ideal_dihedral)
    dihedral_diff = min(delta - actual_dihedral for delta in ideal_dihedrals)
    abs_sin_dihedral = abs(math.sin(math.radians(dihedral_diff)))
    return abs_sin_angle + abs_sin_dihedral

def ideal_bonded_positions(atom, element, geometry=None):
    if geometry is None:
        try:
            geometry = chimera.idatm.typeInfo[atom.idatmType].geometry
        except KeyError:
            geometry = 3

    bond_length = chimera.Element.bondLength(atom.element, element)
    neighbors_crd = [a.xformCoord() for a in atom.neighbors]
    return chimera.bondGeom.bondPositions(atom.xformCoord(), geometry, bond_length, neighbors_crd)


# Copyright 2022 The TEMPO Collaboration
#
# 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.
"""
Module for system 'control operations' as discussed in [Pollock2018].

**[Pollock2018]**
F.  A.  Pollock,  C.  Rodriguez-Rosario,  T.  Frauenheim,
M. Paternostro, and K. Modi, *Non-Markovian quantumprocesses: Complete
framework and efficient characterization*, Phys. Rev. A 97, 012127 (2018).
"""

from typing import Callable, List, Optional, Text, Tuple, Union
from copy import deepcopy

import numpy as np
from numpy import ndarray

from oqupy.base_api import BaseAPIClass
from oqupy.config import NpDtype

class Control(BaseAPIClass):
    """
    Represents a set of system control operations.

    A control operation is a superoperator that acts on the system
    instantaneously at a particular time, as described in [Pollock2018].

    Parameters
    ----------
    dimension: int
        The Hilbert space dimension of the system.
    name: str
        An optional name for the set of control operations.
    description: str
        An optional description of the set of control operations.
    """
    def __init__(
            self,
            dimension: int,
            name: Optional[Text] = None,
            description: Optional[Text] = None) -> None:
        """Creates a Control object. """
        self._dimension = dimension
        self._step_controls = {'pre':{}, 'post':{}}
        self._time_controls = {'pre':{}, 'post':{}}
        self._control_times = {'pre':np.array([]), 'post':np.array([])}
        super().__init__(name, description)

    @property
    def dimension(self):
        """Hilbert space dimension of the controlled system. """
        return self._dimension

    def add_single(
            self,
            time: Union[int, float],
            control_operation: ndarray,
            post: Optional[bool] = False) -> None:
        r"""
        Adds a single control operation at time `time`.

        Parameters
        ----------
        time: Union[int, float]
            The time at which the operation should be applied. If `type(time)`
            is `int` then `time` is understood as the *timestep* to which it
            shall be applied.
        control_operation: ndarray
            The control operation super operator of shape
            :math:`d^2 \times d^2`, where :math:`d` is the system Hilbert space
            dimension.
        post: bool
            If `True` (`False`) the operator is applied at the corresponding
            time step *after* (*before*) a possible measurement of the state.
        """
        if post:
            pre_post = 'post'
        else:
            pre_post = 'pre'

        if isinstance(time, int):
            steps = self._step_controls[pre_post].keys()
            if time in steps:
                self._step_controls[pre_post][time] = \
                    control_operation @ self._step_controls[pre_post][time]
            else:
                self._step_controls[pre_post][time] = control_operation
        elif isinstance(time, float):
            if time in self._control_times[pre_post]:
                self._time_controls[pre_post][time] = \
                    control_operation @ self._time_controls[pre_post][time]
            else:
                self._time_controls[pre_post][time] = control_operation
                times = np.append(self._control_times[pre_post], time)
                times.sort()
                self._control_times[pre_post] = times
        else:
            raise TypeError("Parameter `time` must be either int or float.")

    def add_continuous(
            self,
            control_fct: Callable[[ndarray, float], ndarray],
            post: Optional[bool] = False) -> None:
        """
        ToDo
        """
        raise NotImplementedError()

    def get_controls(
            self,
            step: int,
            dt: Optional[float] = None,
            start_time: Optional[float] = 0.0,
            ) -> Tuple[ndarray, ndarray]:
        """
        Get the pre and post measurement control operation for a specific
        time step.

        Parameters
        ----------
        step: int
            The time step.
        dt: float
            The time step length.
        start_time: float
            The initial time step off-set.

        Returns
        -------
        pre: ndarray
            The control superoperator that should be applied before a state
            measurement.
        post: ndarray
            The control superoperator that should be applied after a state
            measurement.
        """
        pre_control_bool = False
        post_control_bool = False
        pre_control = np.identity(self.dimension**2)
        post_control = np.identity(self.dimension**2)

        # -- pre time-stamp controls --
        a = np.round((self._control_times['pre'] - start_time) / dt)
        times = np.array(self._control_times['pre'])[np.nonzero(a==step)]
        if len(times) > 0:
            print(times)
            pre_control_bool = True
            pre_control = self._time_controls['pre'][times[0]] @ pre_control
            for t in times[1:]:
                pre_control = self._time_controls['pre'][t] @ pre_control

        # -- pre step controls --
        steps = self._step_controls['pre'].keys()
        if step in steps:
            pre_control_bool = True
            pre_control = self._step_controls['pre'][step] @ pre_control

        # -- post step controls --
        steps = self._step_controls['post'].keys()
        if step in steps:
            post_control_bool = True
            post_control = self._step_controls['post'][step] @ post_control

        # -- post time-stamp controls --
        a = np.round((self._control_times['post'] - start_time) / dt)
        times = np.array(self._control_times['post'])[np.nonzero(a==step)]
        if len(times) > 0:
            post_control_bool = True
            post_control = self._time_controls['post'][times[0]] @ post_control
            for t in times[1:]:
                post_control = self._time_controls['post'][t] @ post_control

        if not pre_control_bool:
            pre_control = None
        if not post_control_bool:
            post_control = None
        return pre_control, post_control

class ChainControl(BaseAPIClass):
    """
    Control operations on a linear system chain.

    Parameters
    ----------
    hilbert_space_dimensions: List[int]
        Hilbert space dimension for each chain site.
    name: str
        An optional name for the chain controls.
    description: str
        An optional description of the chain controls.
    """
    def __init__(
            self,
            hilbert_space_dimensions: List[int],
            name: Optional[Text] = None,
            description: Optional[Text] = None) -> None:
        """Create a ChainControl object. """
        tmp_hs_dims = np.array(hilbert_space_dimensions, int)
        assert len(tmp_hs_dims.shape) == 1
        assert len(tmp_hs_dims) >= 1
        assert np.all(tmp_hs_dims > 0)
        self._hs_dims = tmp_hs_dims

        self._single_site_controls_pre = []
        self._single_site_controls_post = []

        super().__init__(name, description)

    def __len__(self):
        """Length of the chain. """
        return len(self._hs_dims)

    @property
    def hs_dims(self):
        """Hilbert space dimensions. """
        return self._hs_dims

    def add_single_site_control(
            self,
            control: ndarray,
            site: int,
            step: int,
            post: Optional[bool] = False,
            name: Optional[Text] = None) -> None:
        """
        Add a control operation at site `site` and time step `step`.

        Parameters
        ----------
        control: ndarray
            Control operation in Liouville space.
        site: int
            Site index.
        step: int
            Timestep to which the control should be applied.
        post: bool
            True if the control should be applied *after* the measurement of
            this time step.
        name: Text
            An optional name to recognize a control operation.
        """
        assert isinstance(site, int)
        assert site < len(self)
        assert isinstance(step, int)
        contr = np.array(control,  dtype=NpDtype)
        assert contr.shape == (self._hs_dims[site]**2, self._hs_dims[site]**2)
        if not post:
            self._single_site_controls_pre.append({
                "contr":contr,
                "site":site,
                "step":step,
                "name":name})
        else:
            self._single_site_controls_post.append({
                "contr":contr,
                "site":site,
                "step":step,
                "name":name})

    def get_single_site_controls(
            self,
            step: int,
            post: bool) -> List[ndarray]:
        """
        Get a list of single site controls for the time step `step`.

        Parameters
        ----------
        step: int
            The time step.
        post: bool
            If `True` (`False`) the set of control superoperators that should be
            applied *after* (*before*) the measurement is returned.

        Returns
        -------
        superoperators_list: list[ndarray]
            List of single site control superoperators.
        """
        empty = True
        controls = [None] * len(self)

        if not post:
            ss_controls = self._single_site_controls_pre
        else:
            ss_controls = self._single_site_controls_post

        for ssc in ss_controls:
            if ssc["step"] == step:
                empty = False
                if controls[ssc["site"]] is None:
                    controls[ssc["site"]] = ssc["contr"]
                else:
                    controls[ssc["site"]] = \
                        controls[ssc["site"]] @ ssc["contr"]

        if empty:
            return None
        return deepcopy(controls)


# pylint: skip-file
import os
import re
import sys
import warnings

xla_flags = os.getenv("XLA_FLAGS", "").lstrip("--")
xla_flags = re.sub(r"xla_force_host_platform_device_count=.+\s", "", xla_flags).split()
os.environ["XLA_FLAGS"] = " ".join([f"--xla_force_host_platform_device_count={100}"])

import aesara.tensor as at
import arviz as az
import jax
import numpy as np
import pandas as pd

from aesara.compile import SharedVariable
from aesara.graph.basic import Apply, Constant, clone, graph_inputs
from aesara.graph.fg import FunctionGraph
from aesara.graph.op import Op
from aesara.graph.opt import MergeOptimizer
from aesara.link.jax.dispatch import jax_funcify
from aesara.tensor.type import TensorType

from pymc3 import modelcontext
from pymc3.aesaraf import compile_rv_inplace

warnings.warn("This module is experimental.")

class NumPyroNUTS(Op):
    def __init__(
        self,
        inputs,
        outputs,
        target_accept=0.8,
        draws=1000,
        tune=1000,
        chains=4,
        seed=None,
        progress_bar=True,
    ):
        self.draws = draws
        self.tune = tune
        self.chains = chains
        self.target_accept = target_accept
        self.progress_bar = progress_bar
        self.seed = seed

        self.inputs, self.outputs = clone(inputs, outputs, copy_inputs=False)
        self.inputs_type = tuple(input.type for input in inputs)
        self.outputs_type = tuple(output.type for output in outputs)
        self.nin = len(inputs)
        self.nout = len(outputs)
        self.nshared = len([v for v in inputs if isinstance(v, SharedVariable)])
        self.samples_bcast = [self.chains == 1, self.draws == 1]

        self.fgraph = FunctionGraph(self.inputs, self.outputs, clone=False)
        MergeOptimizer().optimize(self.fgraph)

        super().__init__()

    def make_node(self, *inputs):

        # The samples for each variable
        outputs = [
            TensorType(v.dtype, self.samples_bcast + list(v.broadcastable))() for v in inputs
        ]

        # The leapfrog statistics
        outputs += [TensorType("int64", self.samples_bcast)()]

        all_inputs = list(inputs)
        if self.nshared > 0:
            all_inputs += self.inputs[-self.nshared :]

        return Apply(self, all_inputs, outputs)

    def do_constant_folding(self, *args):
        return False

    def perform(self, node, inputs, outputs):
        raise NotImplementedError()

@jax_funcify.register(NumPyroNUTS)
def jax_funcify_NumPyroNUTS(op, node, **kwargs):
    from numpyro.infer import MCMC, NUTS

    draws = op.draws
    tune = op.tune
    chains = op.chains
    target_accept = op.target_accept
    progress_bar = op.progress_bar
    seed = op.seed

    # Compile the "inner" log-likelihood function.  This will have extra shared
    # variable inputs as the last arguments
    logp_fn = jax_funcify(op.fgraph, **kwargs)

    if isinstance(logp_fn, (list, tuple)):
        # This handles the new JAX backend, which always returns a tuple
        logp_fn = logp_fn[0]

    def _sample(*inputs):

        if op.nshared > 0:
            current_state = inputs[: -op.nshared]
            shared_inputs = tuple(op.fgraph.inputs[-op.nshared :])
        else:
            current_state = inputs
            shared_inputs = ()

        def log_fn_wrap(x):
            res = logp_fn(
                *(
                    x
                    # We manually obtain the shared values and added them
                    # as arguments to our compiled "inner" function
                    + tuple(
                        v.get_value(borrow=True, return_internal_type=True) for v in shared_inputs
                    )
                )
            )

            if isinstance(res, (list, tuple)):
                # This handles the new JAX backend, which always returns a tuple
                res = res[0]

            return -res

        nuts_kernel = NUTS(
            potential_fn=log_fn_wrap,
            target_accept_prob=target_accept,
            adapt_step_size=True,
            adapt_mass_matrix=True,
            dense_mass=False,
        )

        pmap_numpyro = MCMC(
            nuts_kernel,
            num_warmup=tune,
            num_samples=draws,
            num_chains=chains,
            postprocess_fn=None,
            chain_method="parallel",
            progress_bar=progress_bar,
        )

        pmap_numpyro.run(seed, init_params=current_state, extra_fields=("num_steps",))
        samples = pmap_numpyro.get_samples(group_by_chain=True)
        leapfrogs_taken = pmap_numpyro.get_extra_fields(group_by_chain=True)["num_steps"]
        return tuple(samples) + (leapfrogs_taken,)

    return _sample

def sample_numpyro_nuts(
    draws=1000,
    tune=1000,
    chains=4,
    target_accept=0.8,
    random_seed=10,
    model=None,
    progress_bar=True,
    keep_untransformed=False,
):
    model = modelcontext(model)

    seed = jax.random.PRNGKey(random_seed)

    rv_names = [rv.name for rv in model.value_vars]
    init_state = [model.initial_point[rv_name] for rv_name in rv_names]
    init_state_batched = jax.tree_map(lambda x: np.repeat(x[None, ...], chains, axis=0), init_state)
    init_state_batched_at = [at.as_tensor(v) for v in init_state_batched]

    nuts_inputs = sorted(
        (v for v in graph_inputs([model.logpt]) if not isinstance(v, Constant)),
        key=lambda x: isinstance(x, SharedVariable),
    )
    map_seed = jax.random.split(seed, chains)
    numpyro_samples = NumPyroNUTS(
        nuts_inputs,
        [model.logpt],
        target_accept=target_accept,
        draws=draws,
        tune=tune,
        chains=chains,
        seed=map_seed,
        progress_bar=progress_bar,
    )(*init_state_batched_at)

    # Un-transform the transformed variables in JAX
    sample_outputs = []
    for i, (value_var, rv_samples) in enumerate(zip(model.value_vars, numpyro_samples[:-1])):
        rv = model.values_to_rvs[value_var]
        transform = getattr(value_var.tag, "transform", None)
        if transform is not None:
            untrans_value_var = transform.backward(rv, rv_samples)
            untrans_value_var.name = rv.name
            sample_outputs.append(untrans_value_var)

            if keep_untransformed:
                rv_samples.name = value_var.name
                sample_outputs.append(rv_samples)
        else:
            rv_samples.name = rv.name
            sample_outputs.append(rv_samples)

    print("Compiling...", file=sys.stdout)

    tic1 = pd.Timestamp.now()
    _sample = compile_rv_inplace(
        [],
        sample_outputs + [numpyro_samples[-1]],
        allow_input_downcast=True,
        on_unused_input="ignore",
        accept_inplace=True,
        mode="JAX",
    )
    tic2 = pd.Timestamp.now()

    print("Compilation time = ", tic2 - tic1, file=sys.stdout)

    print("Sampling...", file=sys.stdout)

    *mcmc_samples, leapfrogs_taken = _sample()
    tic3 = pd.Timestamp.now()

    print("Sampling time = ", tic3 - tic2, file=sys.stdout)

    posterior = {k.name: v for k, v in zip(sample_outputs, mcmc_samples)}

    az_trace = az.from_dict(posterior=posterior)

    return az_trace

import numpy,from numpy,from pymc3,import jax


from __future__ import annotations
import qimpy as qp
import numpy as np
import torch
import pathlib
import re
from ._ions_projectors import _get_projectors
from ._ions_atomic import get_atomic_orbitals, get_atomic_density
from ._ions_update import update, _collect_ps_matrix
from typing import Optional, Union, List

class Ions(qp.TreeNode):
    """Ionic system: ionic geometry and pseudopotentials."""

    __slots__ = (
        "n_ions",
        "n_types",
        "symbols",
        "n_ions_type",
        "slices",
        "pseudopotentials",
        "positions",
        "types",
        "M_initial",
        "Z",
        "Z_tot",
        "rho_tilde",
        "Vloc_tilde",
        "n_core_tilde",
        "beta",
        "beta_full",
        "beta_version",
        "D_all",
    )
    n_ions: int  #: number of ions
    n_types: int  #: number of distinct ion types
    n_ions_type: List[int]  #: number of ions of each type
    symbols: List[str]  #: symbol for each ion type
    slices: List[slice]  #: slice to get each ion type
    pseudopotentials: List[qp.ions.Pseudopotential]  #: pseudopotential for each type
    positions: torch.Tensor  #: fractional positions of each ion (n_ions x 3)
    types: torch.Tensor  #: type of each ion (n_ions, int)
    M_initial: Optional[torch.Tensor]  #: initial magnetic moment for each ion
    Z: torch.Tensor  #: charge of each ion type (n_types, float)
    Z_tot: float  #: total ionic charge
    rho_tilde: qp.grid.FieldH  #: ionic charge density (uses coulomb.ion_width)
    Vloc_tilde: qp.grid.FieldH  #: local potential due to ions (including from rho)
    n_core_tilde: qp.grid.FieldH  #: partial core electronic density (for XC)
    beta: qp.electrons.Wavefunction  #: pseudopotential projectors (split-basis only)
    beta_full: Optional[qp.electrons.Wavefunction]  #: full-basis version of `beta`
    beta_version: int  #: version of `beta` to invalidate cached projections
    D_all: torch.Tensor  #: nonlocal pseudopotential matrix (all atoms)

    _get_projectors = _get_projectors
    get_atomic_orbitals = get_atomic_orbitals
    get_atomic_density = get_atomic_density
    update = update
    _collect_ps_matrix = _collect_ps_matrix

    def __init__(
        self,
        *,
        process_grid: qp.utils.ProcessGrid,
        checkpoint_in: qp.utils.CpPath = qp.utils.CpPath(),
        coordinates: Optional[List] = None,
        pseudopotentials: Optional[Union[str, List[str]]] = None,
    ) -> None:
        """Initialize geometry and pseudopotentials.

        Parameters
        ----------
        coordinates
            :yaml:`List of [symbol, x, y, z, args] for each ion in unit cell.`
            Here, symbol is the chemical symbol of the element,
            x, y and z are in the selected coordinate system.
            Optional args is a dictionary of additional per-ion
            parameters that may include:

                * `M`: initial magnetic moment for the ion, which would be a
                  single Mz or [Mx, My, Mz] depending on if the calculation
                  is spinorial. Only specify in spin-polarized calculations.
                * Relaxation constraints: TODO

            Ions of the same type (symbol) must be specified consecutively.
        pseudopotentials
            :yaml:`Pseudopotential filenames or filename templates.`
            Templates for families of pseudopotentials are specified by
            including a $ID in the name which is replaced by the chemical
            symbol of the element. The list of specified file names and
            templates is processed in order, and the first match for
            each element takes precedence.
        """
        super().__init__()
        qp.log.info("\n--- Initializing Ions ---")

        # Read ionic coordinates:
        if coordinates is None:
            coordinates = []
        assert isinstance(coordinates, list)
        self.n_ions = 0  # number of ions
        self.n_types = 0  # number of distinct ion types
        self.symbols = []  # symbol for each ion type
        self.n_ions_type = []  # numebr of ions of each type
        self.slices = []  # slice to get each ion type
        positions = []  # position of each ion
        types = []  # type of each ion (index into symbols)
        M_initial = []  # initial magnetic moments
        type_start = 0
        for coord in coordinates:
            # Check for optional attributes:
            if len(coord) == 4:
                attrib = {}
            elif len(coord) == 5:
                attrib = coord[4]
                if not isinstance(attrib, dict):
                    raise ValueError("ion attributes must be a dict")
            else:
                raise ValueError("each ion must be 4 entries + optional dict")
            # Add new symbol or append to existing:
            symbol = str(coord[0])
            if (not self.symbols) or (symbol != self.symbols[-1]):
                self.symbols.append(symbol)
                self.n_types += 1
                if type_start != self.n_ions:
                    self.slices.append(slice(type_start, self.n_ions))
                    self.n_ions_type.append(self.n_ions - type_start)
                    type_start = self.n_ions
            # Add type and position of current ion:
            types.append(self.n_types - 1)
            positions.append([float(x) for x in coord[1:4]])
            M_initial.append(attrib.get("M", None))
            self.n_ions += 1
        if type_start != self.n_ions:
            self.slices.append(slice(type_start, self.n_ions))  # for last type
            self.n_ions_type.append(self.n_ions - type_start)

        # Check order:
        if len(set(self.symbols)) < self.n_types:
            raise ValueError("coordinates must group ions of same type together")

        # Convert to tensors before storing in class object:
        self.positions = torch.tensor(positions, device=qp.rc.device)
        self.types = torch.tensor(types, device=qp.rc.device, dtype=torch.long)
        # --- Fill in missing magnetizations (if any specified):
        M_lengths = set(
            [(len(M) if isinstance(M, list) else 1) for M in M_initial if M]
        )
        if len(M_lengths) > 1:
            raise ValueError("All M must be same type: 3-vector or scalar")
        elif len(M_lengths) == 1:
            M_length = next(iter(M_lengths))
            assert (M_length == 1) or (M_length == 3)
            M_default = [0.0, 0.0, 0.0] if (M_length == 3) else 0.0
            self.M_initial = torch.tensor(
                [(M if M else M_default) for M in M_initial],
                device=qp.rc.device,
                dtype=torch.double,
            )
        else:
            self.M_initial = None
        self.report()

        # Initialize pseudopotentials:
        self.pseudopotentials = []
        if pseudopotentials is None:
            pseudopotentials = []
        if isinstance(pseudopotentials, str):
            pseudopotentials = [pseudopotentials]
        for i_type, symbol in enumerate(self.symbols):
            fname = None  # full filename for this ion type
            symbol_variants = [symbol.lower(), symbol.upper(), symbol.capitalize()]
            # Check each filename provided in order:
            for ps_name in pseudopotentials:
                if ps_name.count("$ID"):
                    # wildcard syntax
                    for symbol_variant in symbol_variants:
                        fname_test = ps_name.replace("$ID", symbol_variant)
                        if pathlib.Path(fname_test).exists():
                            fname = fname_test  # found
                            break
                else:
                    # specific filename
                    basename = pathlib.PurePath(ps_name).stem
                    ps_symbol = re.split(r"[_\-\.]+", basename)[0]
                    if ps_symbol in symbol_variants:
                        fname = ps_name
                        if not pathlib.Path(fname).exists():
                            raise FileNotFoundError(fname)
                        break
                if fname:
                    break
            # Read pseudopotential file:
            if fname:
                self.pseudopotentials.append(qp.ions.Pseudopotential(fname))
            else:
                raise ValueError(f"no pseudopotential found for {symbol}")
        self.beta_version = 0

        # Calculate total ionic charge (needed for number of electrons):
        self.Z = torch.tensor(
            [ps.Z for ps in self.pseudopotentials], device=qp.rc.device
        )
        self.Z_tot = self.Z[self.types].sum().item()
        qp.log.info(f"\nTotal ion charge, Z_tot: {self.Z_tot:g}")

        # Initialize / check replica process grid dimension:
        n_replicas = 1  # this will eventually change for NEB / phonon DFPT
        process_grid.provide_n_tasks("r", n_replicas)

    def report(self) -> None:
        """Report ionic positions and attributes"""
        qp.log.info(f"{self.n_ions} total ions of {self.n_types} types;" " positions:")
        # Fetch to CPU for reporting:
        positions = self.positions.to(qp.rc.cpu).numpy()
        types = self.types.to(qp.rc.cpu).numpy()
        M_initial = (
            None if (self.M_initial is None) else self.M_initial.to(qp.rc.cpu).numpy()
        )
        for i_ion, position in enumerate(positions):
            # Generate attribute string:
            attrib_str = ""
            attribs = {}
            if M_initial is not None:
                M_i = M_initial[i_ion]
                if np.linalg.norm(M_i):
                    attribs["M"] = M_i
            if attribs:
                attrib_str = ", " + str(attribs).replace("'", "").replace(
                    "array(", ""
                ).replace(")", "")
            # Report:
            qp.log.info(
                f"- [{self.symbols[types[i_ion]]}, {position[0]:11.8f},"
                f" {position[1]:11.8f}, {position[2]:11.8f}{attrib_str}]"
            )

    def translation_phase(
        self, iG: torch.Tensor, atom_slice: slice = slice(None)
    ) -> torch.Tensor:
        """Get translation phases at `iG` for a slice of atoms.
        The result has atoms as the final dimension; summing over that
        dimension yields the structure factor corresponding to these atoms.
        """
        return qp.utils.cis((-2 * np.pi) * (iG @ self.positions[atom_slice].T))

    @property
    def n_projectors(self) -> int:
        """Total number of pseudopotential projectors."""
        return sum(
            (ps.n_projectors * self.n_ions_type[i_ps])
            for i_ps, ps in enumerate(self.pseudopotentials)
        )

    @property
    def n_orbital_projectors(self) -> int:
        """Total number of projectors used to generate atomic orbitals."""
        return sum(
            (ps.n_orbital_projectors * self.n_ions_type[i_ps])
            for i_ps, ps in enumerate(self.pseudopotentials)
        )

    def n_atomic_orbitals(self, n_spinor: int) -> int:
        """Total number of atomic orbitals. This depends on the number
        of spinorial components `n_spinor`."""
        return sum(
            (ps.n_atomic_orbitals(n_spinor) * self.n_ions_type[i_ps])
            for i_ps, ps in enumerate(self.pseudopotentials)
        )