import os
import signal
import sys
import h5py
import lmfit
import numpy as np
import scipy.ndimage as snd
from scipy.spatial.transform import Rotation
import skimage.morphology as skm
import kosselui
from PyQt5.QtCore import QTimer
from PyQt5.QtWidgets import QApplication, QMainWindow, QFileDialog, QListWidgetItem, QMessageBox

import numexpr as ne

axes = [
    [-1, -1, -1],
    [-1, 1, 1],
    [1, -1, 1],
    [1, 1, -1],
    [-1, -1, 1],
    [-1, 1, -1],
    [1, -1, -1],
    [1, 1, 1],
    [-2, -2, 0],
    [-2, 0, -2],
    [-2, 0, 2],
    [-2, 2, 0],
    [0, -2, -2],
    [0, -2, 2],
    [0, 2, -2],
    [0, 2, 2],
    [2, -2, 0],
    [2, 0, -2],
    [2, 0, 2],
    [2, 2, 0],
    [-3, -1, 1],
    [-3, 1, -1],
    [-1, -3, 1],
    [-1, -1, 3],
    [-1, 1, -3],
    [-1, 3, -1],
    [1, -3, -1],
    [1, -1, -3],
    [1, 1, 3],
    [1, 3, 1],
    [3, -1, -1],
    [3, 1, 1],
    [-3, -1, -1],
    [-3, 1, 1],
    [-1, -3, -1],
    [-1, -1, -3],
    [-1, 1, 3],
    [-1, 3, 1],
    [1, -3, 1],
    [1, -1, 3],
    [1, 1, -3],
    [1, 3, -1],
    [3, -1, 1],
    [3, 1, -1],
    [-4, 0, 0],
    [0, -4, 0],
    [0, 0, -4],
    [0, 0, 4],
    [0, 4, 0],
    [4, 0, 0],
    [-4, -2, -2],
    [-4, -2, 2],
    [-4, 2, -2],
    [-4, 2, 2],
    [-2, -4, -2],
    [-2, -4, 2],
    [-2, -2, -4],
    [-2, -2, 4],
    [-2, 2, -4],
    [-2, 2, 4],
    [-2, 4, -2],
    [-2, 4, 2],
    [2, -4, -2],
    [2, -4, 2],
    [2, -2, -4],
    [2, -2, 4],
    [2, 2, -4],
    [2, 2, 4],
    [2, 4, -2],
    [2, 4, 2],
    [4, -2, -2],
    [4, -2, 2],
    [4, 2, -2],
    [4, 2, 2],
    [-3, -3, -1],
    [-3, -1, -3],
    [-3, 1, 3],
    [-3, 3, 1],
    [-1, -3, -3],
    [-1, 3, 3],
    [1, -3, 3],
    [1, 3, -3],
    [3, -3, 1],
    [3, -1, 3],
    [3, 1, -3],
    [3, 3, -1],
    [-3, -3, 1],
    [-3, -1, 3],
    [-3, 1, -3],
    [-3, 3, -1],
    [-1, -3, 3],
    [-1, 3, -3],
    [1, -3, -3],
    [1, 3, 3],
    [3, -3, -1],
    [3, -1, -3],
    [3, 1, 3],
    [3, 3, 1],
    [-4, -4, 0],
    [-4, 0, -4],
    [-4, 0, 4],
    [-4, 4, 0],
    [0, -4, -4],
    [0, -4, 4],
    [0, 4, -4],
    [0, 4, 4],
    [4, -4, 0],
    [4, 0, -4],
    [4, 0, 4],
    [4, 4, 0],
    [-5, -1, -1],
    [-1, -5, -1],
    [-1, -1, -5],
    [-1, 1, 5],
    [-1, 5, 1],
    [1, -1, 5],
    [1, 1, -5],
    [5, -1, 1],
    [5, 1, -1],
    [-5, 1, 1],
    [-3, 3, -3],
    [1, -5, 1],
    [1, 5, -1],
    [3, -3, -3],
    [3, 3, 3],
    [-3, -3, 3],
    [-5, -1, 1],
    [-5, 1, -1],
    [-1, -1, 5],
    [-1, 1, -5],
    [1, -5, -1],
    [1, -1, -5],
    [1, 1, 5],
    [1, 5, 1],
    [5, 1, 1],
    [-3, -3, -3],
    [-3, 3, 3],
    [-1, -5, 1],
    [-1, 5, -1],
    [3, -3, 3],
    [5, -1, -1],
    [3, 3, -3],
    [-6, -2, 0],
    [-6, 0, -2],
    [-6, 0, 2],
    [-6, 2, 0],
    [-2, -6, 0],
    [-2, 0, -6],
    [-2, 0, 6],
    [-2, 6, 0],
    [0, -6, -2],
    [0, -6, 2],
    [0, -2, -6],
    [0, -2, 6],
    [0, 2, -6],
    [0, 2, 6],
    [0, 6, -2],
    [0, 6, 2],
    [2, -6, 0],
    [2, 0, -6],
    [2, 0, 6],
    [2, 6, 0],
    [6, -2, 0],
    [6, 0, -2],
    [6, 0, 2],
    [6, 2, 0],
    [-4, -4, -4],
    [-4, -4, 4],
    [-4, 4, -4],
    [-4, 4, 4],
    [4, -4, -4],
    [4, -4, 4],
    [4, 4, -4],
    [4, 4, 4],
    [-5, -3, 1],
    [-5, -1, 3],
    [-5, 1, -3],
    [-5, 3, -1],
    [-3, -5, 1],
    [-3, -1, 5],
    [-3, 1, -5],
    [-3, 5, -1],
    [-1, -5, 3],
    [-1, -3, 5],
    [-1, 3, -5],
    [-1, 5, -3],
    [1, -5, -3],
    [1, -3, -5],
    [1, 3, 5],
    [1, 5, 3],
    [3, -5, -1],
    [3, -1, -5],
    [3, 1, 5],
    [3, 5, 1],
    [5, -3, -1],
    [5, -1, -3],
    [5, 1, 3],
    [5, 3, 1],
    [-5, -3, -1],
    [-5, -1, -3],
    [-5, 1, 3],
    [-5, 3, 1],
    [-3, -5, -1],
    [-3, -1, -5],
    [-3, 1, 5],
    [-3, 5, 1],
    [-1, -5, -3],
    [-1, -3, -5],
    [-1, 3, 5],
    [-1, 5, 3],
    [1, -5, 3],
    [1, -3, 5],
    [1, 3, -5],
    [1, 5, -3],
    [3, -5, 1],
    [3, -1, 5],
    [3, 1, -5],
    [3, 5, -1],
    [5, -3, 1],
    [5, -1, 3],
    [5, 1, -3],
    [5, 3, -1],
    [-6, -4, -2],
    [-6, -4, 2],
    [-6, -2, -4],
    [-6, -2, 4],
    [-6, 2, -4],
    [-6, 2, 4],
    [-6, 4, -2],
    [-6, 4, 2],
    [-4, -6, -2],
    [-4, -6, 2],
    [-4, -2, -6],
    [-4, -2, 6],
    [-4, 2, -6],
    [-4, 2, 6],
    [-4, 6, -2],
    [-4, 6, 2],
    [-2, -6, -4],
    [-2, -6, 4],
    [-2, -4, -6],
    [-2, -4, 6],
    [-2, 4, -6],
    [-2, 4, 6],
    [-2, 6, -4],
    [-2, 6, 4],
    [2, -6, -4],
    [2, -6, 4],
    [2, -4, -6],
    [2, -4, 6],
    [2, 4, -6],
    [2, 4, 6],
    [2, 6, -4],
    [2, 6, 4],
    [4, -6, -2],
    [4, -6, 2],
    [4, -2, -6],
    [4, -2, 6],
    [4, 2, -6],
    [4, 2, 6],
    [4, 6, -2],
    [4, 6, 2],
    [6, -4, -2],
    [6, -4, 2],
    [6, -2, -4],
    [6, -2, 4],
    [6, 2, -4],
    [6, 2, 4],
    [6, 4, -2],
    [6, 4, 2],
    [-8, 0, 0],
    [0, -8, 0],
    [0, 0, -8],
    [0, 0, 8],
    [0, 8, 0],
    [8, 0, 0],
    [-5, -3, -3],
    [-5, 3, 3],
    [-3, -5, -3],
    [-3, -3, -5],
    [-3, 3, 5],
    [-3, 5, 3],
    [3, -5, 3],
    [3, -3, 5],
    [3, 3, -5],
    [3, 5, -3],
    [5, -3, 3],
    [5, 3, -3],
    [-5, -3, 3],
    [-5, 3, -3],
    [-3, -5, 3],
    [-3, -3, 5],
    [-3, 3, -5],
    [-3, 5, -3],
    [3, -5, -3],
    [3, -3, -5],
    [3, 3, 5],
    [3, 5, 3],
    [5, -3, -3],
    [5, 3, 3],
    [-7, -1, 1],
    [-7, 1, -1],
    [-5, -1, -5],
    [-1, -7, 1],
    [-1, -5, -5],
    [-1, -1, 7],
    [-1, 1, -7],
    [-1, 7, -1],
    [1, -7, -1],
    [1, -1, -7],
    [1, 1, 7],
    [1, 7, 1],
    [7, 1, 1],
    [-5, -5, -1],
    [-5, 1, 5],
    [-5, 5, 1],
    [-1, 5, 5],
    [1, -5, 5],
    [1, 5, -5],
    [5, -5, 1],
    [5, -1, 5],
    [5, 1, -5],
    [5, 5, -1],
    [7, -1, -1],
    [-7, -1, -1],
    [-1, -7, -1],
    [-1, -1, -7],
    [-1, 1, 7],
    [-1, 7, 1],
    [1, -7, 1],
    [1, -1, 7],
    [1, 1, -7],
    [1, 5, 5],
    [1, 7, -1],
    [5, 1, 5],
    [7, -1, 1],
    [7, 1, -1],
    [-7, 1, 1],
    [-5, -5, 1],
    [-5, -1, 5],
    [-5, 1, -5],
    [-5, 5, -1],
    [-1, -5, 5],
    [-1, 5, -5],
    [1, -5, -5],
    [5, -5, -1],
    [5, -1, -5],
    [5, 5, 1],
    [7, 3, 3],
    [-7, -3, 3],
    [-7, 3, -3],
    [-3, -7, 3],
    [-3, -3, 7],
    [-3, 7, -3],
    [-3, 3, -7],
    [3, 7, 3],
    [3, 3, 7],
    [7, -3, -3],
    [3, -7, -3],
    [3, -3, -7],
    [-7, -3, -3],
    [-7, 3, 3],
    [-3, -7, -3],
    [-3, -3, -7],
    [-3, 7, 3],
    [-3, 3, 7],
    [3, 3, -7],
    [3, 7, -3],
    [7, 3, -3],
    [7, -3, 3],
    [3, -7, 3],
    [3, -3, 7],
    [-7, -3, -1],
    [-7, -1, -3],
    [-7, 1, 3],
    [-7, 3, 1],
    [-5, -5, 3],
    [-5, -3, 5],
    [-5, 3, -5],
    [-5, 5, -3],
    [-3, -7, -1],
    [-3, -5, 5],
    [-3, -1, -7],
    [-3, 1, 7],
    [-3, 5, -5],
    [-3, 7, 1],
    [-1, -7, -3],
    [-1, -3, -7],
    [-1, 3, 7],
    [-1, 7, 3],
    [1, -7, 3],
    [1, -3, 7],
    [1, 3, -7],
    [1, 7, -3],
    [3, -7, 1],
    [3, -5, -5],
    [3, -1, 7],
    [3, 1, -7],
    [3, 5, 5],
    [3, 7, -1],
    [5, -5, -3],
    [5, -3, -5],
    [5, 3, 5],
    [5, 5, 3],
    [7, -3, 1],
    [7, -1, 3],
    [7, 1, -3],
    [7, 3, -1],
    [-7, -3, 1],
    [-7, -1, 3],
    [-7, 1, -3],
    [-7, 3, -1],
    [-5, -5, -3],
    [-5, -3, -5],
    [-5, 3, 5],
    [-5, 5, 3],
    [-3, -7, 1],
    [-3, -5, -5],
    [-3, -1, 7],
    [-3, 1, -7],
    [-3, 5, 5],
    [-3, 7, -1],
    [-1, -7, 3],
    [-1, -3, 7],
    [-1, 3, -7],
    [-1, 7, -3],
    [1, -7, -3],
    [1, -3, -7],
    [1, 3, 7],
    [1, 7, 3],
    [3, -7, -1],
    [3, -5, 5],
    [3, -1, -7],
    [3, 1, 7],
    [3, 5, -5],
    [3, 7, 1],
    [5, -5, 3],
    [5, -3, 5],
    [5, 3, -5],
    [5, 5, -3],
    [7, -3, -1],
    [7, -1, -3],
    [7, 1, 3],
    [7, 3, 1],
    [-2, 0, 0],
    [0, -2, 0],
    [0, 0, -2],
    [0, 0, 2],
    [0, 2, 0],
    [2, 0, 0],
    [-2, -2, -2],
    [-2, -2, 2],
    [-2, 2, -2],
    [-2, 2, 2],
    [2, -2, -2],
    [2, -2, 2],
    [2, 2, -2],
    [2, 2, 2],
    [-4, -2, 0],
    [-4, 0, -2],
    [-4, 0, 2],
    [-4, 2, 0],
    [-2, -4, 0],
    [-2, 0, -4],
    [-2, 0, 4],
    [-2, 4, 0],
    [0, -4, -2],
    [0, -4, 2],
    [0, -2, -4],
    [0, -2, 4],
    [0, 2, -4],
    [0, 2, 4],
    [0, 4, -2],
    [0, 4, 2],
    [2, -4, 0],
    [2, 0, -4],
    [2, 0, 4],
    [2, 4, 0],
    [4, -2, 0],
    [4, 0, -2],
    [4, 0, 2],
    [4, 2, 0],
    [-4, -2, -4],
    [-2, -4, -4],
    [2, 4, 4],
    [4, 2, 4],
    [-6, 0, 0],
    [-4, -4, -2],
    [-4, -4, 2],
    [-4, -2, 4],
    [-4, 2, -4],
    [-4, 2, 4],
    [-4, 4, -2],
    [-4, 4, 2],
    [-2, -4, 4],
    [-2, 4, -4],
    [-2, 4, 4],
    [0, -6, 0],
    [0, 0, -6],
    [0, 0, 6],
    [0, 6, 0],
    [2, -4, -4],
    [2, -4, 4],
    [2, 4, -4],
    [4, -4, -2],
    [4, -4, 2],
    [4, -2, -4],
    [4, -2, 4],
    [4, 2, -4],
    [4, 4, -2],
    [4, 4, 2],
    [6, 0, 0],
    [-6, -2, -2],
    [-6, -2, 2],
    [-6, 2, -2],
    [-6, 2, 2],
    [-2, -6, -2],
    [-2, -6, 2],
    [-2, -2, -6],
    [-2, -2, 6],
    [-2, 2, -6],
    [-2, 2, 6],
    [-2, 6, -2],
    [-2, 6, 2],
    [2, -6, -2],
    [2, -6, 2],
    [2, -2, -6],
    [2, -2, 6],
    [2, 2, -6],
    [2, 2, 6],
    [2, 6, -2],
    [2, 6, 2],
    [6, -2, -2],
    [6, -2, 2],
    [6, 2, -2],
    [6, 2, 2],
    [-6, -4, 0],
    [-6, 0, -4],
    [-4, -6, 0],
    [-4, 0, -6],
    [-4, 0, 6],
    [-4, 6, 0],
    [0, -6, -4],
    [0, -4, -6],
    [0, -4, 6],
    [0, 4, -6],
    [0, 4, 6],
    [0, 6, 4],
    [4, -6, 0],
    [4, 0, -6],
    [4, 0, 6],
    [4, 6, 0],
    [6, 0, 4],
    [6, 4, 0],
    [-6, 0, 4],
    [-6, 4, 0],
    [0, -6, 4],
    [0, 6, -4],
    [6, -4, 0],
    [6, 0, -4],
]

axesactive = tuple(
    [
        list(active)
        for active in set(
            [  # GaAs 1
                (-1, 1, 1),
                (-1, -1, 1),
                (1, -1, 1),
                (-2, -2, 0),
                (2, -2, 0),
                (2, 2, 0),
                (-2, 2, 0),
                (3, 1, 1),
                (-3, -1, 1),
                (-1, -3, 1),
                (2, 2, 4),
                (-2, 2, 4),
                (-2, -2, 4),
                (2, -2, 4),
                (0, -4, 4),
                (0, 4, 4),
                (4, 0, 4),
                (-4, 0, 4),
            ]
            + [  # GaAs 2
                (0, 2, 2),
                (2, -2, 0),
                (-2, 0, 2),
                (2, 2, 0),
                (0, -2, 2),
                (-2, 2, 0),
                (2, 0, 2),
                (-2, -2, 0),
                (0, 4, 4),
                (-4, 0, 4),
                (4, 0, 4),
                (0, -4, 4),
                (-1, 3, 3),
                (3, 1, 3),
                (-1, -3, 3),
                (1, 3, 3),
                (-3, -1, 3),
                (1, -3, 3),
                (2, 0, 6),
                (0, -2, 6),
                (0, 2, 6),
                (-2, 0, 6),
            ]
        )
    ]
)

def getpoints(axis, testpoints, E, a, N):
    normaxis = np.linalg.norm(axis)
    return np.where(
        ne.evaluate(
            "abs(tp1*a1+tp2*a2+tp3*a3-c)<0.001",
            {
                "c": (12398 * normaxis ** 2 / (E * 2 * a)),
                "tp1": testpoints[:, 0],
                "tp2": testpoints[:, 1],
                "tp3": testpoints[:, 2],
                "a1": axis[0],
                "a2": axis[1],
                "a3": axis[2],
            },
        ).reshape(N)
    )

def residual(params, testpointsdet, axes):
    r = Rotation.from_euler("xyz", np.array((params["rot_x"].value, params["rot_y"].value, params["rot_z"].value)))
    testpoints = r.apply(testpointsdet + np.array((params["trans_x"], params["trans_y"], params["trans_z"])), inverse=False)
    testpoints = testpoints / np.linalg.norm(testpoints, axis=1)[:, None]
    res = (np.dot(testpoints, axes.T)) / (np.linalg.norm(axes, axis=1)) ** 2 - params["c"].value
    res = np.take_along_axis(res, np.nanargmin(np.abs(res), axis=1)[:, None], 1)
    return res

class Kossel(QMainWindow, kosselui.Ui_MainWindow):
    def __init__(self):
        super(self.__class__, self).__init__()
        self.N0 = 1000
        self.N1 = 1000

        self.setupUi(self)
        self.plotButton.clicked.connect(self.plot_data)
        self.clearButton.clicked.connect(self.clear_data)
        self.loadButton.clicked.connect(self.load_file)
        self.peakButton.clicked.connect(self.find_peaks)

        ## Not implemented
        self.peakButton.setEnabled(False)
        self.thresholdSlider.setEnabled(False)

        self.saveButton.clicked.connect(self.save)
        self.fitButton.clicked.connect(self.fit_data)
        self.datasetCombo.currentTextChanged.connect(self.plot_bg)
        self.removeBackgroundBox.stateChanged.connect(self.plot_bg)
        self.data = np.zeros((self.N0, self.N1))
        self.data[0, 0] = 1
        self.data[-1, -1] = 1
        self.inputfile = None
        self.bgplot = self.plotarea.canvas.ax.matshow(self.data, vmin=0, vmax=1)
        self.peaks = np.zeros((self.N0, self.N1))
        self.peaks[:] = np.nan
        self.peaksplot = self.plotarea.canvas.ax.matshow(self.peaks, vmin=0, vmax=1, cmap="gray")
        self.plotarea.canvas.ax.set_xlim(0, self.N1)
        self.plotarea.canvas.ax.set_ylim(self.N0, 0)
        self.plotarea.canvas.draw_idle()
        self.rangeSlider.startValueChanged.connect(self.setclim)
        self.rangeSlider.endValueChanged.connect(self.setclim)
        self.redrawtimer = QTimer()
        self.redrawtimer.setSingleShot(True)
        self.redrawtimer.timeout.connect(self.plotarea.canvas.draw)
        for i, ax in enumerate(axes):
            label = np.array2string(np.array(ax), precision=0)
            item = QListWidgetItem(label)
            item.setData(1, ax)
            self.reflexList.addItem(item)
            if ax in axesactive:
                item.setSelected(True)
        for el in [
            self.angleXSpin,
            self.angleYSpin,
            self.angleZSpin,
            self.transXSpin,
            self.transYSpin,
            self.transZSpin,
            self.energySpin,
            self.latticeSpin,
        ]:
            el.valueChanged.connect(self.invalidate_plot)

        self.plotpoints = {}
        self.testpoints = None

        def styles(i=0):
            colors = (
                "#1f77b4",
                "#aec7e8",
                "#ff7f0e",
                "#ffbb78",
                "#2ca02c",
                "#98df8a",
                "#d62728",
                "#ff9896",
                "#9467bd",
                "#c5b0d5",
                "#8c564b",
                "#c49c94",
                "#e377c2",
                "#f7b6d2",
                "#7f7f7f",
                "#c7c7c7",
                "#bcbd22",
                "#dbdb8d",
                "#17becf",
                "#9edae5",
            )
            markers = (".", "+", "x")
            while True:
                yield colors[i % len(colors)], markers[(i // len(colors)) % len(markers)]
                i += 1

        self.sgen = styles()
        self.inputfilename = None

        self.plotarea.canvas.mpl_connect("motion_notify_event", self.mpl_move)
        self.plotarea.canvas.mpl_connect("button_release_event", self.mpl_release)

    def mpl_move(self, event):
        if event.button == 1:
            tmp = self.data[int(event.ydata) - 1 : int(event.ydata) + 2, int(event.xdata) - 1 : int(event.xdata) + 2]
            m = tmp == tmp.max()
            self.peaks[int(event.ydata) - 1 : int(event.ydata) + 2, int(event.xdata) - 1 : int(event.xdata) + 2][m] = 1
        elif event.button == 3:
            self.peaks[int(event.ydata) - 5 : int(event.ydata) + 5, int(event.xdata) - 5 : int(event.xdata) + 5] = np.nan

    def mpl_release(self, event):
        self.peaksplot.set_array(self.peaks)
        self.plotarea.canvas.draw_idle()

    def plot_bg(self):
        if self.inputfile is not None:
            data = np.array(self.inputfile[self.datasetCombo.currentText()])
            if data.ndim == 2:
                if self.removeBackgroundBox.isChecked():
                    data = data - snd.grey_opening(data, structure=skm.disk(10))
                data = data - np.nanmin(data)
                data = data / np.nanmax(data)
                self.data = data
                self.bgplot.remove()
                self.bgplot = self.plotarea.canvas.ax.matshow(self.data, vmin=0, vmax=1, zorder=0)
                if data.shape[0] != self.N0 or data.shape[1] != self.N1:
                    self.N0 = data.shape[0]
                    self.N1 = data.shape[1]
                    self.peaks = np.zeros((self.N0, self.N1))
                    self.peaks[:] = np.nan
                    self.peaksplot = self.plotarea.canvas.ax.matshow(self.peaks, vmin=0, vmax=1, cmap="gray", zorder=1)
                    self.invalidate_plot()
                    self.plotarea.canvas.draw_idle()

                self.plotarea.canvas.ax.set_xlim(0, self.N1)
                self.plotarea.canvas.ax.set_ylim(self.N0, 0)

        self.plotarea.canvas.draw_idle()

    def save(self):
        filename, _ = QFileDialog.getSaveFileName(self, "save plot", "", "pdf (*.pdf)")
        if filename:
            try:
                self.plotarea.canvas.ax.get_figure().savefig(filename)
            except Exception:
                msg = QMessageBox()
                msg.setIcon(QMessageBox.Critical)
                msg.setText("Saving failed")
                msg.setWindowTitle("Error")
                retval = msg.exec_()

    def setclim(self):
        clim = np.array(self.rangeSlider.getRange()) / 100
        self.bgplot.set_clim(clim)
        self.redrawtimer.start(100)
        # self.plotarea.canvas.draw_idle()

    def invalidate_plot(self):
        for label in self.plotpoints:
            self.plotpoints[label][0] = None
        self.testpoints = None

    def load_file(self):
        fname = QFileDialog.getOpenFileName(self, "Open file", ".", "h5 (*.h5 *.hdf5)")
        self.inputfileLabel.setText(os.path.split(fname[0])[1])
        self.inputfilename = fname[0]
        try:
            if self.inputfile is not None:
                self.inputfile.close()
                self.inputfile = None
            inputfile = h5py.File(fname[0], "r")
            self.datasetCombo.clear()
            self.datasetCombo.addItems(inputfile.keys())
            self.inputfile = inputfile
            self.plot_bg()
        except Exception as e:
            print("error opening file:", e)

    def fit_data(self):
        x, y = np.where(~np.isnan(self.peaks))
        testpointsdet = np.array((x, y, [0] * len(x))).T - np.array((self.N0 // 2, self.N1 // 2, 0))
        fit_params = lmfit.Parameters()
        fit_params.add(
            "rot_x",
            value=self.angleXSpin.value() / 180 * np.pi,
            min=(self.angleXSpin.value() - 15) / 180 * np.pi,
            max=(self.angleXSpin.value() + 15) / 180 * np.pi,
            vary=not self.rotXFixBox.isChecked(),
        )
        fit_params.add(
            "rot_y",
            value=self.angleYSpin.value() / 180 * np.pi,
            min=(self.angleYSpin.value() - 15) / 180 * np.pi,
            max=(self.angleYSpin.value() + 15) / 180 * np.pi,
            vary=not self.rotYFixBox.isChecked(),
        )
        fit_params.add(
            "rot_z",
            value=self.angleZSpin.value() / 180 * np.pi,
            min=(self.angleZSpin.value() - 15) / 180 * np.pi,
            max=(self.angleZSpin.value() + 15) / 180 * np.pi,
            vary=not self.rotZFixBox.isChecked(),
        )
        fit_params.add(
            "trans_x",
            value=self.transXSpin.value(),
            min=self.transXSpin.value() - 100,
            max=self.transXSpin.value() + 100,
            vary=not self.transXFixBox.isChecked(),
        )
        fit_params.add(
            "trans_y",
            value=self.transYSpin.value(),
            min=self.transYSpin.value() - 100,
            max=self.transYSpin.value() + 100,
            vary=not self.transYFixBox.isChecked(),
        )
        fit_params.add(
            "trans_z",
            value=self.transZSpin.value(),
            min=self.transZSpin.value() - 100,
            max=self.transZSpin.value() + 100,
            vary=not self.transZFixBox.isChecked(),
        )
        c = 12.398 / (self.energySpin.value() * 2 * self.latticeSpin.value())
        fit_params.add("c", value=c, min=0.8 * c, max=1.2 * c, vary=not self.latticeFixBox.isChecked())
        axs = np.array([np.array(item.data(1)) for item in self.reflexList.selectedItems()])
        minner = lmfit.Minimizer(residual, fit_params, fcn_args=(testpointsdet, axs))
        result = minner.minimize(method="bfgs")
        print(lmfit.fit_report(result))
        self.angleXSpin.setValue(result.params["rot_x"] * 180 / np.pi)
        self.angleYSpin.setValue(result.params["rot_y"] * 180 / np.pi)
        self.angleZSpin.setValue(result.params["rot_z"] * 180 / np.pi)
        self.transXSpin.setValue(result.params["trans_x"])
        self.transYSpin.setValue(result.params["trans_y"])
        self.transZSpin.setValue(result.params["trans_z"])
        self.latticeSpin.setValue(12.398 / (self.energySpin.value() * 2 * result.params["c"]))

    def plot_data(self):
        if self.testpoints is None:
            self.clear_data()
            r = Rotation.from_euler("xyz", np.array((self.angleXSpin.value(), self.angleYSpin.value(), self.angleZSpin.value())) / 180 * np.pi)
            Y, X, Z = np.meshgrid(np.arange(-self.N1 // 2, self.N1 // 2, 1), np.arange(-self.N0 // 2, self.N0 // 2, 1), 0)
            testpoints = np.array([m.ravel() for m in [X, Y, Z]]).T
            testpoints = r.apply(testpoints + np.array((self.transXSpin.value(), self.transYSpin.value(), self.transZSpin.value())), inverse=False)
            self.testpoints = testpoints / np.linalg.norm(testpoints, axis=1)[:, None]

        E = self.energySpin.value() * 1000
        a = self.latticeSpin.value()

        items = self.reflexList.selectedItems()
        selectedlabels = [item.data(0) for item in items]
        for label, item in self.plotpoints.items():
            if label not in selectedlabels:
                if item[1] is not None:
                    item[1].remove()
                    item[1] = None
                if item[2] is not None:
                    item[2].remove()
                    item[2] = None

        for k, item in enumerate(items):
            ax = np.array(item.data(1), dtype=int)
            label = np.array2string(ax, precision=0)
            ax = ax.astype(float)
            self.progressBar.setValue(int(k / len(items) * 100))
            self.progressBar.update()
            QApplication.processEvents()
            if label in self.plotpoints and self.plotpoints[label][0] is not None:
                continue
            points = getpoints(ax, testpoints=self.testpoints, E=E, a=a, N=(self.N0, self.N1))
            self.plotpoints[label] = [points, None, None]

        for label in selectedlabels:
            points = self.plotpoints[label][0]

            if len(points[0]) > 0:
                s = next(self.sgen)
                if self.plotpoints[label][1] is None:
                    self.plotpoints[label][1] = self.plotarea.canvas.ax.scatter(points[1], points[0], label=label, c=s[0], s=1)

                if self.plotpoints[label][2] is None:
                    for j in range(15):
                        i = np.random.choice(np.arange(0, len(points[0])))
                        if 10 < points[0][i] < (self.N0 - 20) and 10 < points[1][i] < (self.N1 - 100):
                            self.plotpoints[label][2] = self.plotarea.canvas.ax.text(points[1][i] + 5, points[0][i] + 5, s=label, c=s[0])
                            break
                    else:
                        self.plotpoints[label][2] = self.plotarea.canvas.ax.text(
                            np.clip(points[1][i] + 5, 20, self.N1 - 20), np.clip(points[0][i] + 5, 20, self.N0 - 100), s=label, c=s[0]
                        )

        self.plotarea.canvas.ax.set_xlim(0, self.N1)
        self.plotarea.canvas.ax.set_ylim(self.N0, 0)
        self.plotarea.canvas.draw_idle()

    def clear_data(self):
        for label in self.plotpoints:
            if self.plotpoints[label][1] is not None:
                self.plotpoints[label][1].remove()
                self.plotpoints[label][1] = None

            if self.plotpoints[label][2] is not None:
                self.plotpoints[label][2].remove()
                self.plotpoints[label][2] = None

        self.plotarea.canvas.draw_idle()

    def find_peaks(self):
        if self.data is not None:
            self.peaks[:] = np.nan
            self.peaks[self.data > (self.thresholdSlider.value() / 100)] = 1
            self.peaksplot.set_array(self.peaks)
            self.plotarea.canvas.draw_idle()

def sigint_handler(*args):
    sys.stderr.write("\r")
    QApplication.quit()

if __name__ == "__main__":
    print('starting kossel')
    signal.signal(signal.SIGINT, sigint_handler)
    app = QApplication(sys.argv)
    timer = QTimer()
    timer.start(250)
    timer.timeout.connect(lambda: None)
    form = Kossel()
    form.show()
    r = app.exec_()
    sys.exit(r)

import numpy,import scipy,from scipy,import numexpr


#!/usr/bin/python
# -*- coding: utf-8 -*-
import numpy as np
import re

from .libpyfeos import  _get_calc_limits, _initialize, _init_mat,\
       _get_eos_all,_get_eos, _get_mat_par, _get_crit_point,\
       _get_soft_shere_par, _get_energy_offsets, _delete_material

from ..base import _pull_tables, TableBase, MaterialBase, GridBase
from copy import deepcopy

avalable_tabs = np.unique([tab.format(s=spec) \
                for tab in ['P{s}_DT', 'U{s}_DT', 'S{s}_DT', 'A{s}_DT', 'Zfc_DT',
                            'P{s}_DU{s}', 'T_DU{s}', 'S{s}_DU{s}', 'A{s}_DU{s}', 'Zfc_DU{s}']\
                for spec in ['t', 'iz', 'ec', 'e', 'ic']])

# define a global global_table_handles variable
if 'global_table_handles' not in globals():
    global_table_handles = None

class FeosTable(TableBase):
    _original_units = 'feos'

    def __init__(self, _name, table_handle=None, options={},info={}, units='cgs'):
        self._id = table_handle
        self._name = _name
        self._Nelements  = info['Nelements']
        self.update(info)
        super(FeosTable, self).__init__(_name, table_handle, options, units)
        self.options = options

    def _interpolate(self, X, Y, kind):
        if type(X) in [float, int]:
           X = np.array(X)
           Y = np.array(Y)
        if X.shape != Y.shape:
            raise ValueError('X and Y arguments should be ndarrays of the same shape!')

        init_shape = X.shape

        Xvals = np.array(X.flatten(), dtype='float64')*self._X_convert
        Yvals = np.array(Y.flatten(), dtype='float64')*self._Y_convert
        Fvar, XY_vars =  self._name.split('_')
        if XY_vars == 'DT':
            if kind == 'F':
                res = _get_eos_all(self._id, self._Nelements,
                                          self.options['use_maxwell'], Xvals, Yvals)
                fVals_f = res[Fvar]*self._F_convert

                fVals = fVals_f.reshape(init_shape)
                return fVals
            else:
                raise NotImplemented
        else:
            raise NotImplemented

    def __getitem__(self, key):
        """ Overwriting default dict's __getitem__ method """
        if key in ['abar', "Mean_Atomic_Mass"]:
            return self['Atot']/self['Xtot']
        if key in ['zbar', "Mean_Atomic_Num"]:
            return self['Ztot']/self['Xtot']
        if key == "Normal_Density": key = "rho_ref"
        if key == "Modulus": key = "bulk_mod_ref"
        if key == "R_Array": key = "D_Array"
        if key == "Exchange_Coeff":
            return 0.0
        return super(FeosTable, self).__getitem__(key)

class FeosMaterial(MaterialBase):

    _default_options = {'use_maxwell': True, 'use_softspheres': True,
            'maxwell_temp_arr': None, 'max_materials': 1, 'debug': False,
            'rho_grid': None, 'temp_grid': None,
            'grid_subsample_default': 0, 'grid_kind': 'solid',
            'interpolation_mode_inversed': 'linear', 'precalculate': True}
    _backend = 'feos'
    _original_units = 'feos'

    def __init__(self, material=None, tables=['Pt_DT', 'Ut_DT'],
        options={'use_maxwell': True, 'use_softspheres': True,
            'maxwell_temp_arr': None, 'max_materials': 1, 'debug': False},
            spec=['t'], units='cgs'):
        """
        Parameters:
        -----------
         - material: int: 4 digit SESAME material ID
         - tables: list: ['table1_id', 'table2_id', ..etc]
         - options: dict: {'tableid_regexpr': {'opt1': optval}, etc}
            For example {".t_DT": {'create_tzero': True}} would apply the
            EOS_CREATE_TZERO option to both 'Pt_DT' and 'Ut_DT' tables. The
            tableid_regexpr accepts regular expressions.
            [see  help(re) for more details].

        """
        self.options = self._validate_options(options)
        opt = self.options

        self.tables = _pull_tables(tables, spec, avalable_tabs)
        self.material = int(material)

        # making a global array to store table handles
        # This makes sure that FEOS is only initalized once
        global global_table_handles
        if global_table_handles is None:
            self._id = 1
            _initialize(opt['max_materials'], int(opt['debug']))
            global_table_handles = np.zeros(opt['max_materials'])
        else:
            for idx0, mat in enumerate(global_table_handles):
                if not mat:
                    self._id = idx0 + 1
                    break
            else:
                raise ValueError('Could not allocate {0} material.\n\
                    Please reinitialize all materials and increase the "max_materials" option!'.format(
                    self.material))

        global_table_handles[self._id-1] = self.material

        self.info = self._get_info_init()

        # default grid
        default_grid = GridBase(self.info['rho_ref'], kind=opt['grid_kind'],
                subsample_temp=opt['grid_subsample_default'],
                subsample_rho=opt['grid_subsample_default'])

        grid_units = self._set_units(units, 'Pt_DT')
        if opt['rho_grid'] is None:
            rho_grid = default_grid.rho_grid*grid_units.o2r('D', 'cgs', units)
        else:
            rho_grid = opt['rho_grid']

        if opt['temp_grid'] is None:
            temp_grid = default_grid.temp_grid*grid_units.o2r('T', 'cgs', units)
        else:
            temp_grid = opt['temp_grid']

        self.info['T_Array'] = temp_grid
        self.info['D_Array'] = rho_grid
        temp_grid_feos = temp_grid*grid_units.o2r('T', units, 'feos')
        rho_grid_feos = rho_grid*grid_units.o2r('D', units, 'feos')

        self.info.update(self._get_info_final(temp_grid_feos))
        if opt['precalculate']:
            self.precalculated = self._precalculate_on_grid(rho_grid_feos,temp_grid_feos)
            self.precalculated['D_Array'] = rho_grid_feos
            self.precalculated['T_Array'] = temp_grid_feos
        else:
            self.precalculated = {}
        for tab_idx, tab_key in enumerate(self.tables):
           this_info = deepcopy(self.info)
           if self.precalculated:
               F_Array = self.precalculated[tab_key.split('_')[0]]*\
                        grid_units.o2r(tab_key[0], 'feos', units)
               this_info['F_Array'] =  F_Array.T

           setattr(self, tab_key,
                FeosTable(tab_key,
                            self._id,
                            options=opt,
                            info=this_info,
                            units=units))
        self._init_base()

    def _precalculate_on_grid(self, rho_grid, temp_grid):
        R, T = np.meshgrid(rho_grid, temp_grid, indexing='ij')
        R_flat, T_flat = R.flatten(), T.flatten()
        #res = self._Nelements
        res_tmp = _get_eos_all(self._id, self._Nelements,
                self.options['use_maxwell'], R_flat, T_flat)
        res = {}

        for key in res_tmp:
            res[key] = res_tmp[key].reshape(R.shape)
        return res

    def _get_info_init(self):
        """ We have to preinitialize the material just to get the rho_ref """
        opt = self.options
        # preinitializint without Maxwell constructions to go faster
        Nelements, N_maxwell_iso, Maxwell_temp_reliable = _init_mat(self._id,
                self.material, False, False,
                np.logspace(-4, 1, 100))

        info  = {}
        info.update(_get_mat_par(self._id, Nelements))
        _delete_material(self._id)

        return info

    def _get_info_final(self, temp_grid):
        """ We have to preinitialize the material just to get the density """
        opt = self.options
        T_mask = temp_grid<100. # eV

        Nelements, N_maxwell_iso, Maxwell_temp_reliable = _init_mat(self._id,
                self.material, opt['use_maxwell'], opt['use_softspheres'],
                temp_grid)

        info  = {}
        self._Nelements = Nelements
        info['Nelements'] = Nelements
        info['N_maxwell_iso'] = N_maxwell_iso
        info['Maxwell_temp_reliable'] = Maxwell_temp_reliable
        info.update(_get_mat_par(self._id, Nelements))
        info.update(_get_energy_offsets(self._id))
        if opt['use_maxwell']:
            info.update(_get_crit_point(self._id))
            if opt['use_softspheres']:
                info.update(_get_soft_shere_par(self._id))

        return info


#!/usr/bin/env python

from __future__ import division, print_function

import sys, os
import copy
import numpy as np
import simtk.openmm as mm
import simtk.openmm.app as app
import simtk.unit as u
import parmed as pmd
from parmed.openmm.reporters import NetCDFReporter
from pymbar import mbar
from openmm_surface_affinities_lib import *
import waterlib as wl
from scipy import optimize

#Given a topology and structure file, this script sets up a simulation of a solvated solute
#(this code is for just in bulk!) and periodically kicks off NVE simulations from the NPT
#configurations and temperatures. These NVE simulations are then used to assess dynamics, while
#the trajectory in the NPT can be used to evaluate solute properties in the fully coupled
#ensemble. The below script applies to bulk systems.

def normalExponential(t, A, Tau):
  #A function to define a normal exponential for fitting decay of water residency in shells
  return A*np.exp(-(t/Tau))

def stretchedExponential(t, A, Tau, B):
  #A function to define a stretched exponential for fitting the dipole vector autocorrelation function
  return A*np.exp(-(t/Tau)**B)

def doSimDynamics(top, systemRef, integratorRef, platform, prop, temperature, scalexy=False, inBulk=False, state=None, pos=None, vels=None, nSteps=10000000):
  #Input a topology object, reference system, integrator, platform, platform properties,
  #and optionally state file, positions, or velocities
  #If state is specified including positions and velocities and pos and vels are not None, the
  #positions and velocities from the provided state will be overwritten
  #Does NPT, stopping periodically to run NVE to compute dynamics
  #Only the NPT simulation will be saved, not the NVE

  #Copy the reference system and integrator objects
  system = copy.deepcopy(systemRef)
  integrator = copy.deepcopy(integratorRef)

  #For NPT, add the barostat as a force
  #If not in bulk, use anisotropic barostat
  if not inBulk:
    system.addForce(mm.MonteCarloAnisotropicBarostat((1.0, 1.0, 1.0)*u.bar,
                                                     temperature, #Temperature should be SAME as for thermostat
                                                     scalexy, #Set with flag for flexibility
                                                     scalexy,
                                                     True, #Only scale in z-direction
                                                     250 #Time-steps between MC moves
                                                    )
    )
  #If in bulk, have to use isotropic barostat to avoid any weird effects with box changing dimensions
  else:
    system.addForce(mm.MonteCarloBarostat(1.0*u.bar,
                                          temperature,
                                          250
                                         )
    )

  #Create new simulation object for NPT simulation
  sim = app.Simulation(top.topology, system, integrator, platform, prop, state)

  #Also create copies and simulation object for the NVE we will be running
  systemNVE = copy.deepcopy(systemRef)
  integratorNVE = mm.VerletIntegrator(2.0*u.femtoseconds)
  integratorNVE.setConstraintTolerance(1.0E-08)
  simNVE = app.Simulation(top.topology, systemNVE, integratorNVE, platform, prop)

  #Set the particle positions in the NPT simulation
  if pos is not None:
    sim.context.setPositions(pos)

  #Apply constraints before starting the simulation
  sim.context.applyConstraints(1.0E-08)

  #Check starting energy decomposition if want
  #decompEnergy(sim.system, sim.context.getState(getPositions=True))

  #Initialize velocities if not specified
  if vels is not None:
    sim.context.setVelocities(vels)
  else:
    try:
      testvel = sim.context.getState(getVelocities=True).getVelocities()
      print("Velocities included in state, starting with 1st particle: %s"%str(testvel[0]))
      #If all the velocities are zero, then set them to the temperature
      if not np.any(testvel.value_in_unit(u.nanometer/u.picosecond)):
        print("Had velocities, but they were all zero, so setting based on temperature.")
        sim.context.setVelocitiesToTemperature(temperature)
    except:
      print("Could not find velocities, setting with temperature")
      sim.context.setVelocitiesToTemperature(temperature)

  #Set up the reporter to output energies, volume, etc.
  sim.reporters.append(app.StateDataReporter(
                                             'prod_out.txt', #Where to write - can be stdout or file name (default .csv, I prefer .txt)
                                             500, #Number of steps between writes
                                             step=True, #Write step number
                                             time=True, #Write simulation time
                                             potentialEnergy=True, #Write potential energy
                                             kineticEnergy=True, #Write kinetic energy
                                             totalEnergy=True, #Write total energy
                                             temperature=True, #Write temperature
                                             volume=True, #Write volume
                                             density=False, #Write density
                                             speed=True, #Estimate of simulation speed
                                             separator='  ' #Default is comma, but can change if want (I like spaces)
                                            )
  )

  #Set up reporter for printing coordinates (trajectory)
  sim.reporters.append(NetCDFReporter(
                                      'prod.nc', #File name to write trajectory to
                                      500, #Number of steps between writes
                                      crds=True, #Write coordinates
                                      vels=True, #Write velocities
                                      frcs=False #Write forces
                                     )
  )

  #Identify solute indices and water oxygen indices
  soluteInds = []
  owInds = []
  hw1Inds = []
  hw2Inds = []
  for res in top.residues:
    if res.name not in ['OTM', 'CTM', 'STM', 'NTM', 'SOL']:
      for atom in res.atoms:
        soluteInds.append(atom.idx)
    elif res.name == 'SOL':
      for atom in res.atoms:
        if atom.name == 'OW':
          owInds.append(atom.idx)
        elif atom.name == 'HW1':
          hw1Inds.append(atom.idx)
        elif atom.name == 'HW2':
          hw2Inds.append(atom.idx)

  print("Solute indices:")
  print(soluteInds)
  #print("Water oxygen indices:")
  #print(owInds)
  #print("Water hydrogen (1st) indices:")
  #print(hw1Inds)
  #print("Water hydrogen (2nd) indices:")
  #print(hw2Inds)

  #Define cutoffs for solute solvation shells
  solShell1Cut = 0.55 #nanometers from all solute atoms (including hydrogens)
  solShell2Cut = 0.85

  #Create array to store the dynamic information of interest every 0.2 ps (100 steps) for 50 ps
  calcSteps = 100
  calcTotSteps = 25000
  numWats = np.zeros((int(calcTotSteps/calcSteps)+1, 2)) #Number waters that started in shell that are in shell at later time
  dipCorrs = np.zeros((int(calcTotSteps/calcSteps)+1, 2)) #Dipole correlation in both solute shells

  #Start running dynamics
  print("\nRunning NPT simulation with interspersed NVE to find dynamics...")
  sim.context.setTime(0.0)

  stepChunk = 5000 #Run NVE for 50 ps to find dynamics every 10 ps
  countSteps = 0

  while countSteps < nSteps:

    countSteps += stepChunk
    sim.step(stepChunk)

    #Record the simulation state so can kick off the NVE simulation
    thisState = sim.context.getState(getPositions=True, getVelocities=True)

    #Get solute and water oxygen coordinates after wrapping around the solute
    coords = thisState.getPositions(asNumpy=True)
    boxDims = np.diagonal(thisState.getPeriodicBoxVectors(asNumpy=True))
    wrapCOM = np.average(coords[soluteInds], axis=0)
    coords = wl.reimage(coords, wrapCOM, boxDims) - wrapCOM
    solCoords = coords[soluteInds]
    owCoords = coords[owInds]
    hw1Coords = coords[hw1Inds]
    hw2Coords = coords[hw2Inds]

    #Figure out which waters are in the solute solvation shells
    shell1BoolMat = wl.nearneighbors(solCoords, owCoords, boxDims, 0.0, solShell1Cut)
    shell1Bool = np.array(np.sum(shell1BoolMat, axis=0), dtype=bool)
    shell2BoolMat = wl.nearneighbors(solCoords, owCoords, boxDims, solShell1Cut, solShell2Cut)
    shell2Bool = np.array(np.sum(shell2BoolMat, axis=0), dtype=bool)

    #Count number of waters in each shell (will need for averaging)
    thisCount1 = int(np.sum(shell1Bool))
    thisCount2 = int(np.sum(shell2Bool))

    #print("Found %i waters in shell1"%thisCount1)
    #print("Found %i waters in shell2"%thisCount2)

    #Loop over waters in shells and compute dipole vectors as references
    refDipoles1 = np.zeros((thisCount1, 3))
    refDipoles2 = np.zeros((thisCount2, 3))
    for k, pos in enumerate(owCoords[shell1Bool]):
      thisOHvecs = wl.reimage([hw1Coords[shell1Bool][k], hw2Coords[shell1Bool][k]], pos, boxDims) - pos
      thisDip = -0.5*(thisOHvecs[0] + thisOHvecs[1])
      refDipoles1[k] = thisDip / np.linalg.norm(thisDip)
    for k, pos in enumerate(owCoords[shell2Bool]):
      thisOHvecs = wl.reimage([hw1Coords[shell2Bool][k], hw2Coords[shell2Bool][k]], pos, boxDims) - pos
      thisDip = -0.5*(thisOHvecs[0] + thisOHvecs[1])
      refDipoles2[k] = thisDip / np.linalg.norm(thisDip)

    #Set up the NVE simulation
    simNVE.context.setState(thisState)
    simNVE.context.setTime(0.0)

    #Loop over taking steps to computed dynamics
    countStepsNVE = 0
    while countStepsNVE <= calcTotSteps:
      calcState = simNVE.context.getState(getPositions=True)
      #Get solute and water oxygen coordinates after wrapping around the solute
      coords = calcState.getPositions(asNumpy=True)
      wrapCOM = np.average(coords[soluteInds], axis=0)
      coords = wl.reimage(coords, wrapCOM, boxDims) - wrapCOM
      solCoords = coords[soluteInds]
      owCoords = coords[owInds]
      hw1Coords = coords[hw1Inds]
      hw2Coords = coords[hw2Inds]
      #Count waters that started in each shell that are now in the shell at this time
      #No absorbing boundaries
      thisbool1Mat = wl.nearneighbors(solCoords, owCoords, boxDims, 0.0, solShell1Cut)
      thisbool1 = np.array(np.sum(thisbool1Mat, axis=0), dtype=bool)
      thisbool2Mat = wl.nearneighbors(solCoords, owCoords, boxDims, solShell1Cut, solShell2Cut)
      thisbool2 = np.array(np.sum(thisbool2Mat, axis=0), dtype=bool)
      numWats[int(countStepsNVE/calcSteps),0] += int(np.sum(thisbool1*shell1Bool))
      numWats[int(countStepsNVE/calcSteps),1] += int(np.sum(thisbool2*shell2Bool))
      #Loop over waters in shells and compute dipole vectors for this configuration
      #Adding to sum that we will normalize to find average at each time point
      for k, pos in enumerate(owCoords[shell1Bool]):
        thisOHvecs = wl.reimage([hw1Coords[shell1Bool][k], hw2Coords[shell1Bool][k]], pos, boxDims) - pos
        thisDip = -0.5*(thisOHvecs[0] + thisOHvecs[1])
        thisDip /= np.linalg.norm(thisDip)
        dipCorrs[int(countStepsNVE/calcSteps),0] += (np.dot(thisDip, refDipoles1[k]) / float(thisCount1))
      for k, pos in enumerate(owCoords[shell2Bool]):
        thisOHvecs = wl.reimage([hw1Coords[shell2Bool][k], hw2Coords[shell2Bool][k]], pos, boxDims) - pos
        thisDip = -0.5*(thisOHvecs[0] + thisOHvecs[1])
        thisDip /= np.linalg.norm(thisDip)
        dipCorrs[int(countStepsNVE/calcSteps),1] += (np.dot(thisDip, refDipoles2[k]) / float(thisCount2))
      simNVE.step(calcSteps)
      countStepsNVE += calcSteps

  #Finish normalizing dipole correlations (really cosine of angle between dipole vector at different times)
  numWats /= float(int(nSteps/stepChunk))
  dipCorrs /= float(int(nSteps/stepChunk))
  print("Normalizing factor for finding averages: %f"%float(int(nSteps/stepChunk)))

  #And save the final state of the NPT simulation in case we want to extend it
  sim.saveState('nptDynamicsState.xml')

  #And return the dipole correlations and times at which they were computed
  timeVals = 0.002*np.arange(0.0, calcTotSteps+0.0001, calcSteps)

  return numWats, dipCorrs, timeVals

def main(args):
  #Get the structure and topology files from the command line
  #ParmEd accepts a wide range of file types (Amber, GROMACS, CHARMM, OpenMM... but not LAMMPS)
  try:
    topFile = args[0]
    strucFile = args[1]
  except IndexError:
    print("Specify topology and structure files from the command line.")
    sys.exit(2)

  print("Using topology file: %s" % topFile)
  print("Using structure file: %s" % strucFile)

  print("\nSetting up system...")

  #Load in the files for initial simulations
  top = pmd.load_file(topFile)
  struc = pmd.load_file(strucFile)

  #Transfer unit cell information to topology object
  top.box = struc.box[:]

  #Set up some global features to use in all simulations
  temperature = 298.15*u.kelvin

  #Define the platform (i.e. hardware and drivers) to use for running the simulation
  #This can be CUDA, OpenCL, CPU, or Reference
  #CUDA is for NVIDIA GPUs
  #OpenCL is for CPUs or GPUs, but must be used for old CPUs (not SSE4.1 compatible)
  #CPU only allows single precision (CUDA and OpenCL allow single, mixed, or double)
  #Reference is a clear, stable reference for other code development and is very slow, using double precision by default
  platform = mm.Platform.getPlatformByName('CUDA')
  prop = {#'Threads': '2', #number of threads for CPU - all definitions must be strings (I think)
          'Precision': 'mixed', #for CUDA or OpenCL, select the precision (single, mixed, or double)
          'DeviceIndex': '0', #selects which GPUs to use - set this to zero if using CUDA_VISIBLE_DEVICES
          'DeterministicForces': 'True' #Makes sure forces with CUDA and PME are deterministic
         }

  #Create the OpenMM system that can be used as a reference
  systemRef = top.createSystem(
                               nonbondedMethod=app.PME, #Uses PME for long-range electrostatics, simple cut-off for LJ
                               nonbondedCutoff=12.0*u.angstroms, #Defines cut-off for non-bonded interactions
                               rigidWater=True, #Use rigid water molecules
                               constraints=app.HBonds, #Constrains all bonds involving hydrogens
                               flexibleConstraints=False, #Whether to include energies for constrained DOFs
                               removeCMMotion=True, #Whether or not to remove COM motion (don't want to if part of system frozen)
  )

  #Set up the integrator to use as a reference
  integratorRef = mm.LangevinIntegrator(
                                        temperature, #Temperature for Langevin
                                        1.0/u.picoseconds, #Friction coefficient
                                        2.0*u.femtoseconds, #Integration timestep
  )
  integratorRef.setConstraintTolerance(1.0E-08)

  #Get solute atoms
  soluteIndices = []
  for res in top.residues:
    if res.name not in ['OTM', 'CTM', 'STM', 'NTM', 'SOL']:
      for atom in res.atoms:
        soluteIndices.append(atom.idx)

  print("\nSolute indices: %s" % str(soluteIndices))

  #JUST for boric acid, add a custom bonded force
  #Couldn't find a nice, compatible force field, but did find A forcefield, so using it
  #But has no angle terms on O-B-O and instead a weird bond repulsion term
  #This term also prevents out of plane bending
  #Simple in our case because boric acid is symmetric, so only need one parameter
  #Parameters come from Otkidach and Pletnev, 2001
  #Here, Ad = (A^2) / (d^6) since Ai and Aj and di and dj are all the same
  #In the original paper, B-OH bond had A = 1.72 and d = 0.354
  #Note that d is dimensionless and A should have units of (Angstrom^3)*(kcal/mol)^(1/2)
  #These units are inferred just to make things work out with kcal/mol and the given distance dependence
  bondRepulsionFunction = 'Ad*(1.0/r)^6'
  BondRepulsionForce = mm.CustomBondForce(bondRepulsionFunction)
  BondRepulsionForce.addPerBondParameter('Ad') #Units are technically kJ/mol * nm^6
  baOxInds = []
  for aind in soluteIndices:
    if top.atoms[aind].type == 'oh':
      baOxInds.append(aind)
  for i in range(len(baOxInds)):
    for j in range(i+1, len(baOxInds)):
      BondRepulsionForce.addBond(baOxInds[i], baOxInds[j], [0.006289686])

  systemRef.addForce(BondRepulsionForce)

  #Setting up the alchemical system so we can repeat the calculation with a decoupled particle
  #We need to add a custom non-bonded force for the solute being alchemically changed
  #Will be helpful to have handle on non-bonded force handling LJ and coulombic interactions
  NBForce = None
  for frc in systemRef.getForces():
    if (isinstance(frc, mm.NonbondedForce)):
      NBForce = frc

  #Turn off dispersion correction since have interface
  NBForce.setUseDispersionCorrection(False)

  forceLabelsRef = getForceLabels(systemRef)

  decompEnergy(systemRef, struc.positions, labels=forceLabelsRef)

  #Separate out alchemical and regular particles using set objects
  alchemicalParticles = set(soluteIndices)
  chemicalParticles = set(range(systemRef.getNumParticles())) - alchemicalParticles

  #Define the soft-core function for turning on/off LJ interactions
  #In energy expressions for CustomNonbondedForce, r is a special variable and refers to the distance between particles
  #All other variables must be defined somewhere in the function.
  #The exception are variables like sigma1 and sigma2.
  #It is understood that a parameter will be added called 'sigma' and that the '1' and '2' are to specify the combining rule.
  softCoreFunction = '4.0*lambdaLJ*epsilon*x*(x-1.0); x = (1.0/reff_sterics);'
  softCoreFunction += 'reff_sterics = (0.5*(1.0-lambdaLJ) + ((r/sigma)^6));'
  softCoreFunction += 'sigma=0.5*(sigma1+sigma2); epsilon = sqrt(epsilon1*epsilon2)'
  #Define the system force for this function and its parameters
  SoftCoreForce = mm.CustomNonbondedForce(softCoreFunction)
  SoftCoreForce.addGlobalParameter('lambdaLJ', 1.0) #Throughout, should follow convention that lambdaLJ=1.0 is fully-interacting state
  SoftCoreForce.addPerParticleParameter('sigma')
  SoftCoreForce.addPerParticleParameter('epsilon')

  #Will turn off electrostatics completely in the original non-bonded force
  #In the end-state, only want electrostatics inside the alchemical molecule
  #To do this, just turn ON a custom force as we turn OFF electrostatics in the original force
  ONE_4PI_EPS0 = 138.935456 #in kJ/mol nm/e^2
  soluteCoulFunction = '(1.0-(lambdaQ^2))*ONE_4PI_EPS0*charge/r;'
  soluteCoulFunction += 'ONE_4PI_EPS0 = %.16e;' % (ONE_4PI_EPS0)
  soluteCoulFunction += 'charge = charge1*charge2'
  SoluteCoulForce = mm.CustomNonbondedForce(soluteCoulFunction)
  #Note this lambdaQ will be different than for soft core (it's also named differently, which is CRITICAL)
  #This lambdaQ corresponds to the lambda that scales the charges to zero
  #To turn on this custom force at the same rate, need to multiply by (1.0-lambdaQ**2), which we do
  SoluteCoulForce.addGlobalParameter('lambdaQ', 1.0)
  SoluteCoulForce.addPerParticleParameter('charge')

  #Also create custom force for intramolecular alchemical LJ interactions
  #Could include with electrostatics, but nice to break up
  #We could also do this with a separate NonbondedForce object, but it would be a little more work, actually
  soluteLJFunction = '4.0*epsilon*x*(x-1.0); x = (sigma/r)^6;'
  soluteLJFunction += 'sigma=0.5*(sigma1+sigma2); epsilon=sqrt(epsilon1*epsilon2)'
  SoluteLJForce = mm.CustomNonbondedForce(soluteLJFunction)
  SoluteLJForce.addPerParticleParameter('sigma')
  SoluteLJForce.addPerParticleParameter('epsilon')

  #Loop over all particles and add to custom forces
  #As we go, will also collect full charges on the solute particles
  #AND we will set up the solute-solute interaction forces
  alchemicalCharges = [[0]]*len(soluteIndices)
  for ind in range(systemRef.getNumParticles()):
    #Get current parameters in non-bonded force
    [charge, sigma, epsilon] = NBForce.getParticleParameters(ind)
    #Make sure that sigma is not set to zero! Fine for some ways of writing LJ energy, but NOT OK for soft-core!
    if sigma/u.nanometer == 0.0:
      newsigma = 0.3*u.nanometer #This 0.3 is what's used by GROMACS as a default value for sc-sigma
    else:
      newsigma = sigma
    #Add the particle to the soft-core force (do for ALL particles)
    SoftCoreForce.addParticle([newsigma, epsilon])
    #Also add the particle to the solute only forces
    SoluteCoulForce.addParticle([charge])
    SoluteLJForce.addParticle([sigma, epsilon])
    #If the particle is in the alchemical molecule, need to set it's LJ interactions to zero in original force
    if ind in soluteIndices:
      NBForce.setParticleParameters(ind, charge, sigma, epsilon*0.0)
      #And keep track of full charge so we can scale it right by lambda
      alchemicalCharges[soluteIndices.index(ind)] = charge

  #Now we need to handle exceptions carefully
  for ind in range(NBForce.getNumExceptions()):
    [p1, p2, excCharge, excSig, excEps] = NBForce.getExceptionParameters(ind)
    #For consistency, must add exclusions where we have exceptions for custom forces
    SoftCoreForce.addExclusion(p1, p2)
    SoluteCoulForce.addExclusion(p1, p2)
    SoluteLJForce.addExclusion(p1, p2)

  #Only compute interactions between the alchemical and other particles for the soft-core force
  SoftCoreForce.addInteractionGroup(alchemicalParticles, chemicalParticles)

  #And only compute alchemical/alchemical interactions for other custom forces
  SoluteCoulForce.addInteractionGroup(alchemicalParticles, alchemicalParticles)
  SoluteLJForce.addInteractionGroup(alchemicalParticles, alchemicalParticles)

  #Set other soft-core parameters as needed
  SoftCoreForce.setCutoffDistance(12.0*u.angstroms)
  SoftCoreForce.setNonbondedMethod(mm.CustomNonbondedForce.CutoffPeriodic)
  SoftCoreForce.setUseLongRangeCorrection(False)
  systemRef.addForce(SoftCoreForce)

  #Set other parameters as needed - note that for the solute force would like to set no cutoff
  #However, OpenMM won't allow a bunch of potentials with cutoffs then one without...
  #So as long as the solute is smaller than the cut-off, won't have any problems!
  SoluteCoulForce.setCutoffDistance(12.0*u.angstroms)
  SoluteCoulForce.setNonbondedMethod(mm.CustomNonbondedForce.CutoffPeriodic)
  SoluteCoulForce.setUseLongRangeCorrection(False)
  systemRef.addForce(SoluteCoulForce)

  SoluteLJForce.setCutoffDistance(12.0*u.angstroms)
  SoluteLJForce.setNonbondedMethod(mm.CustomNonbondedForce.CutoffPeriodic)
  SoluteLJForce.setUseLongRangeCorrection(False)
  systemRef.addForce(SoluteLJForce)

  #First do simulation with fully coupled state
  SoftCoreForce.setGlobalParameterDefaultValue(0, 1.0)
  SoluteCoulForce.setGlobalParameterDefaultValue(0, 1.0)

  for k, ind in enumerate(soluteIndices):
    [charge, sig, eps] = NBForce.getParticleParameters(ind)
    NBForce.setParticleParameters(ind, alchemicalCharges[k]*1.0, sig, eps)

  forceLabelsRef = getForceLabels(systemRef)
  decompEnergy(systemRef, struc.positions, labels=forceLabelsRef)

  os.mkdir('coupled')
  os.chdir('coupled')

  #Do NVT simulation
  stateFileNVT, stateNVT = doSimNVT(top, systemRef, integratorRef, platform, prop, temperature, pos=struc.positions)

  #And do NPT simulation using state information from NVT
  stateFileNPT, stateNPT = doSimNPT(top, systemRef, integratorRef, platform, prop, temperature, inBulk=True, state=stateFileNVT)

  #Now perform dynamics simulation to get dynamics - this is defined here, NOT in openmm_surface_affinities_lib.py
  numShellWaters, dipoleCosAng, timePoints = doSimDynamics(top, systemRef, integratorRef, platform, prop, temperature, inBulk=True, state=stateFileNPT)

  #Finally, want to now save the water residency over time and then also fit to exponential decay
  np.savetxt("shell_watCounts_coupled.txt", np.hstack((np.array([timePoints]).T, numShellWaters)),
             header="Time (ps)  Number waters in the 1st and 2nd solvation shells")

  opt1, pcov1 = optimize.curve_fit(normalExponential, timePoints, numShellWaters[:,0]/numShellWaters[0,0])
  decayTime1 = opt1[1]
  opt2, pcov2 = optimize.curve_fit(normalExponential, timePoints, numShellWaters[:,1]/numShellWaters[0,1])
  decayTime2 = opt2[1]

  print("\nIn the fully coupled ensemble:")
  print("\tWater residency correlation time for 1st shell waters: %f"%decayTime1)
  print("\tWater residency correlation time for 2nd shell waters: %f"%decayTime2)

  #Finally, want to now save the dipoles over time and then also fit to stretched exponential
  np.savetxt("rotational_timeCorr_coupled.txt", np.hstack((np.array([timePoints]).T, dipoleCosAng)),
             header="Time (ps)  Cos(angle) between starting dipole and dipole for 1st and 2nd solvation shells")

  opt1, pcov1 = optimize.curve_fit(stretchedExponential, timePoints, dipoleCosAng[:,0])
  decayTime1 = opt1[1]
  opt2, pcov2 = optimize.curve_fit(stretchedExponential, timePoints, dipoleCosAng[:,1])
  decayTime2 = opt2[1]

  print("\tRotational correlation time for 1st shell waters: %f"%decayTime1)
  print("\tRotational correlation time for 2nd shell waters: %f"%decayTime2)

  os.chdir('../')

  #Next simulate with decoupled state, but do same analysis
  #At least this way the volumes considered will be similar
  SoftCoreForce.setGlobalParameterDefaultValue(0, 0.0)
  SoluteCoulForce.setGlobalParameterDefaultValue(0, 0.0)

  for k, ind in enumerate(soluteIndices):
    [charge, sig, eps] = NBForce.getParticleParameters(ind)
    NBForce.setParticleParameters(ind, alchemicalCharges[k]*0.0, sig, eps)

  forceLabelsRef = getForceLabels(systemRef)
  decompEnergy(systemRef, struc.positions, labels=forceLabelsRef)

  os.mkdir('decoupled')
  os.chdir('decoupled')

  #Do NVT simulation
  stateFileNVT, stateNVT = doSimNVT(top, systemRef, integratorRef, platform, prop, temperature, pos=struc.positions)

  #And do NPT simulation using state information from NVT
  stateFileNPT, stateNPT = doSimNPT(top, systemRef, integratorRef, platform, prop, temperature, inBulk=True, state=stateFileNVT)

  #Now perform dynamics simulation to get dynamics - this is defined here, NOT in openmm_surface_affinities_lib.py
  numShellWaters, dipoleCosAng, timePoints = doSimDynamics(top, systemRef, integratorRef, platform, prop, temperature, inBulk=True, state=stateFileNPT)

  #Finally, want to now save the water residency over time and then also fit to exponential decay
  np.savetxt("shell_watCounts_decoupled.txt", np.hstack((np.array([timePoints]).T, numShellWaters)),
             header="Time (ps)  Number waters in the 1st and 2nd solvation shells")

  opt1, pcov1 = optimize.curve_fit(normalExponential, timePoints, numShellWaters[:,0]/numShellWaters[0,0])
  decayTime1 = opt1[1]
  opt2, pcov2 = optimize.curve_fit(normalExponential, timePoints, numShellWaters[:,1]/numShellWaters[0,1])
  decayTime2 = opt2[1]

  print("\nIn the perfectly decoupled ensemble:")
  print("\tWater residency correlation time for 1st shell waters: %f"%decayTime1)
  print("\tWater residency correlation time for 2nd shell waters: %f"%decayTime2)

  #Finally, want to now save the dipoles over time and then also fit to stretched exponential
  np.savetxt("rotational_timeCorr_decoupled.txt", np.hstack((np.array([timePoints]).T, dipoleCosAng)),
             header="Time (ps)  Cos(angle) between starting dipole and dipole for 1st and 2nd solvation shells")

  opt1, pcov1 = optimize.curve_fit(stretchedExponential, timePoints, dipoleCosAng[:,0])
  decayTime1 = opt1[1]
  opt2, pcov2 = optimize.curve_fit(stretchedExponential, timePoints, dipoleCosAng[:,1])
  decayTime2 = opt2[1]

  print("\tRotational correlation time for 1st shell waters: %f"%decayTime1)
  print("\tRotational correlation time for 2nd shell waters: %f"%decayTime2)

  os.chdir('../')

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


from typing import List

import matplotlib as mpl
import numpy as np
from astropy import units as u
from astropy.coordinates import CartesianRepresentation
from matplotlib import pyplot as plt

from poliastro.plotting.util import BODY_COLORS, generate_label
from poliastro.util import norm

from ._base import Trajectory

class StaticOrbitPlotter:
    """StaticOrbitPlotter class.

    This class holds the perifocal plane of the first
    :py:class:`~poliastro.twobody.orbit.Orbit` plotted in it using
    :py:meth:`plot`, so all following
    plots will be projected on that plane. Alternatively, you can call
    :py:meth:`set_frame` to set the frame before plotting.

    """

    def __init__(self, ax=None, num_points=150, dark=False):
        """Constructor.

        Parameters
        ----------
        ax : ~matplotlib.axes.Axes
            Axes in which to plot. If not given, new ones will be created.
        num_points : int, optional
            Number of points to use in plots, default to 150.
        dark : bool, optional
            If set as True, plots the orbit in Dark mode.
        """
        self.ax = ax
        if not self.ax:
            if dark:
                with plt.style.context("dark_background"):
                    _, self.ax = plt.subplots(figsize=(6, 6))
            else:
                _, self.ax = plt.subplots(figsize=(6, 6))
        self.num_points = num_points
        self._frame = None
        self._attractor = None
        self._attractor_radius = np.inf * u.km
        self._trajectories = []  # type: List[Trajectory]

    @property
    def trajectories(self):
        return self._trajectories

    def set_frame(self, p_vec, q_vec, w_vec):
        """Sets perifocal frame.

        Raises
        ------
        ValueError
            If the vectors are not a set of mutually orthogonal unit vectors.
        """
        if not np.allclose([norm(v) for v in (p_vec, q_vec, w_vec)], 1):
            raise ValueError("Vectors must be unit.")
        elif not np.allclose([p_vec.dot(q_vec), q_vec.dot(w_vec), w_vec.dot(p_vec)], 0):
            raise ValueError("Vectors must be mutually orthogonal.")
        else:
            self._frame = p_vec, q_vec, w_vec

        if self._trajectories:
            self._redraw()

    def _redraw(self):
        for artist in self.ax.lines + self.ax.collections:
            artist.remove()

        for trajectory, state, label, color in self._trajectories:
            self._plot(trajectory, state, label, color)

        self.ax.relim()
        self.ax.autoscale()

    def _plot_trajectory(self, trajectory, color=None):
        rr = trajectory.represent_as(CartesianRepresentation).xyz.transpose()
        x, y = self._project(rr)
        lines = self.ax.plot(x.to(u.km).value, y.to(u.km).value, "--", color=color)

        return lines

    def plot_trajectory(self, trajectory, *, label=None, color=None):
        """Plots a precomputed trajectory.

        Parameters
        ----------
        trajectory : ~astropy.coordinates.BaseRepresentation, ~astropy.coordinates.BaseCoordinateFrame
            Trajectory to plot.
        label : str, optional
            Label.
        color : str, optional
            Color string.

        """
        if self._attractor is None or self._frame is None:
            raise ValueError(
                "An attractor and a frame must be set up first, please use "
                "set_attractor(Major_Body) and set_frame(*orbit.pqw()) "
                "or plot(orbit)."
            )

        self._redraw_attractor(
            trajectory.represent_as(CartesianRepresentation).norm().min() * 0.15
        )  # Arbitrary threshold
        lines = self._plot_trajectory(trajectory, color)

        if label:
            lines[0].set_label(label)
            self.ax.legend(
                loc="upper left", bbox_to_anchor=(1.05, 1.015), title="Names and epochs"
            )

        self._trajectories.append(
            Trajectory(trajectory, None, label, lines[0].get_color())
        )

        return lines

    def set_attractor(self, attractor):
        """Sets plotting attractor.

        Parameters
        ----------
        attractor : ~poliastro.bodies.Body
            Central body.

        """
        if self._attractor is None:
            self._attractor = attractor

        elif attractor is not self._attractor:
            raise NotImplementedError(
                "Attractor has already been set to {}.".format(self._attractor.name)
            )

    def _project(self, rr):
        rr_proj = rr - rr.dot(self._frame[2])[:, None] * self._frame[2]
        x = rr_proj.dot(self._frame[0])
        y = rr_proj.dot(self._frame[1])
        return x, y

    def _redraw_attractor(self, min_radius=0 * u.km):
        radius = max(self._attractor.R.to(u.km), min_radius.to(u.km))
        color = BODY_COLORS.get(self._attractor.name, "#999999")

        for attractor in self.ax.findobj(match=mpl.patches.Circle):
            attractor.remove()

        if radius < self._attractor_radius:
            self._attractor_radius = radius

        self.ax.add_patch(
            mpl.patches.Circle((0, 0), self._attractor_radius.value, lw=0, color=color)
        )

    def _plot(self, trajectory, state=None, label=None, color=None):
        lines = self._plot_trajectory(trajectory, color)

        if state is not None:
            x0, y0 = self._project(state[None])

            # Plot current position
            l, = self.ax.plot(
                x0.to(u.km).value,
                y0.to(u.km).value,
                "o",
                mew=0,
                color=lines[0].get_color(),
            )
            lines.append(l)

        if label:
            if not self.ax.get_legend():
                size = self.ax.figure.get_size_inches() + [8, 0]
                self.ax.figure.set_size_inches(size)

            # This will apply the label to either the point or the osculating
            # orbit depending on the last plotted line
            # NOTE: What about generating both labels,
            # indicating that one is the osculating orbit?
            lines[-1].set_label(label)
            self.ax.legend(
                loc="upper left", bbox_to_anchor=(1.05, 1.015), title="Names and epochs"
            )

        self.ax.set_xlabel("$x$ (km)")
        self.ax.set_ylabel("$y$ (km)")
        self.ax.set_aspect(1)

        return lines

    def plot(self, orbit, label=None, color=None):
        """Plots state and osculating orbit in their plane.
        """
        if not self._frame:
            self.set_frame(*orbit.pqw())

        self.set_attractor(orbit.attractor)
        self._redraw_attractor(orbit.r_p * 0.15)  # Arbitrary threshold
        positions = orbit.sample(self.num_points)
        if label:
            label = generate_label(orbit, label)

        lines = self._plot(positions, orbit.r, label, color)

        self._trajectories.append(
            Trajectory(positions, orbit.r, label, lines[0].get_color())
        )
        return lines


from typing import Dict, List, Union, Optional, Tuple
from warnings import warn
from pathlib import Path
import pandas as pd
import numpy as np
from mendeleev import element
from chemicalc.instruments import InstConfig
from chemicalc.utils import (
    doppler_shift,
    convolve_spec,
    calc_gradient,
)
from chemicalc.file_mgmt import (
    data_dir,
    download_package_files,
    precomputed_res,
    precomputed_ref_id,
    precomputed_label_id,
    precomputed_alpha_included,
)

# noinspection PyTypeChecker
elements_included: List[str] = [x.symbol for x in element(list(range(3, 100)))]
"""
List[str]: List of all elements included in the pre-computed spectral grids.
"""
alpha_el: List[str] = ["O", "Ne", "Mg", "Si", "S", "Ar", "Ca", "Ti"]
"""
List[str]: List of elements considered when calculating bulk [:math:`\\alpha`/H].
"""

class ReferenceSpectra:
    """
    Object for spectra of a specific reference star

    :param str reference: Name of reference star to load (e.g., 'RGB_m1.5')
    :param str init_res: Initial resolution of high-res reference spectra.
                         Only 300000 is presently included for default spectra.
                         Can be approximate if using custom reference spectra.
    :param bool scale_by_iron: If true, scales all elemental abundances by [Fe/H]
    :param bool alpha_included: If true, will include an alpha label after the atmospheric parameters
                                and before the other elements (i.e., between v_micro and Li)
    :param \**kwargs: see below

    :keyword str ref_spec_file: Full path to file of reference spectra
    :keyword str ref_label_file: Full path to file of reference spectra labels

    :ivar Dict[Union[str,float],int] resolution: Dictionary of resolving powers for each instrument
    :ivar ref_spec_file: Full path to file of reference spectra
    :ivar ref_label_file: Full path to file of reference spectra labels
    :ivar Dict[Union[str,float],np.ndarray] wavelength: Dictionary of wavelength arrays for each instrument
    :ivar Dict[Union[str,float],np.ndarray] spectra: Dictionary of spectral grids for each instrument
    :ivar pd.DataFrame labels: Labels corresponding to each spectrum in the grid
    :ivar Dict[Union[str,float],pd.DataFrame] gradients: Dictionary of partial derivatives for each instrument
    :ivar int nspectra: Number of spectra included in the grid
    :ivar int nlabels: Number of labels included
    """

    def __init__(
        self,
        reference: str,
        init_res: float = precomputed_res[0],
        scale_by_iron: bool = False,
        alpha_included: bool = True,
        **kwargs,
    ) -> None:
        if not isinstance(reference, str):
            raise TypeError("reference must be str")
        if not isinstance(init_res, (int, float)):
            raise TypeError("init_res must be float")
        self.reference = reference
        self.resolution = {"init": init_res}

        if "ref_spec_file" in kwargs:
            self.ref_spec_file = Path(kwargs["ref_spec_file"])
            if not self.ref_spec_file.exists():
                raise ValueError(f"ref_spec_file {self.ref_spec_file} does not exist")
        else:
            if not self.resolution["init"] in precomputed_res:
                raise ValueError(f"{init_res} not a precomputed resolution")
            self.ref_spec_file = data_dir.joinpath(
                f"reference_spectra_{init_res:06}.h5"
            )
            if not self.ref_spec_file.exists():
                print(
                    "Downloading reference file---this may take a few minutes but is only necessary once"
                )
                download_package_files(
                    id_str=precomputed_ref_id[init_res], destination=self.ref_spec_file
                )

        if "ref_label_file" in kwargs:
            self.ref_label_file = Path(kwargs["ref_label_file"])
            if not self.ref_label_file.exists():
                raise ValueError(f"ref_label_file {self.ref_label_file} does not exist")
        else:
            self.ref_label_file = data_dir.joinpath("reference_labels.h5")
            if not self.ref_label_file.exists():
                print(
                    "Downloading label_file---this should be quick and is only necessary once"
                )
                download_package_files(
                    id_str=precomputed_label_id, destination=self.ref_label_file
                )
            if alpha_included and reference not in precomputed_alpha_included:
                raise ValueError(
                    f"alpha offsets not currently included for {reference}"
                )

        ref_list_spec = list(
            pd.DataFrame(pd.read_hdf(self.ref_spec_file, "ref_list")).values.flatten()
        )
        ref_list_label = list(
            pd.DataFrame(pd.read_hdf(self.ref_label_file, "ref_list")).values.flatten()
        )
        if not (reference in ref_list_spec) and (reference in ref_list_label):
            raise ValueError(
                f"{reference} is not included in ref_label_file and/or ref_spec_file"
            )

        wave_df = pd.DataFrame(pd.read_hdf(self.ref_spec_file, "highres_wavelength"))
        spec_df = pd.DataFrame(pd.read_hdf(self.ref_spec_file, reference))
        label_df = pd.DataFrame(pd.read_hdf(self.ref_label_file, reference))
        if scale_by_iron:
            label_df.loc[set(elements_included) ^ {"Fe"}] -= label_df.loc["Fe"]
        if alpha_included:
            label_df = pd.concat(
                [
                    label_df.iloc[:3],
                    pd.DataFrame(label_df.loc[alpha_el].mean()).T,
                    label_df.iloc[3:],
                ]
            )
            label_df.index = ["Teff", "logg", "v_micro", "alpha"] + elements_included
            if (
                np.abs(
                    label_df.loc["alpha"][[4, 6, 7]].max() - label_df.loc["alpha"][0]
                )
                < 0.001
            ):
                warn(
                    "Expected offset in alpha not found. Are you sure this reference spectra includes alpha gradients?"
                    + "\nIf so, they must come immediately after v_micro offsets in both the label and spectra files.",
                    UserWarning,
                )
        else:
            label_df.index = ["Teff", "logg", "v_micro"] + elements_included

        self.wavelength = dict(init=wave_df.to_numpy().T[0])
        self.spectra = dict(init=spec_df.to_numpy().T)
        self.labels = label_df
        self.gradients: Dict[Union[str, float], pd.DataFrame] = {}

        self.nspectra = self.spectra["init"].shape[0]
        self.nlabels = self.labels.shape[0]

    def add_rv_spec(self, d_rv: float, symmetric: bool = True) -> None:
        """
        Adds spectra and labels corresponding to a small doppler shift of the reference spectra.
        Assumes that the first spectra in ref_spec_file is a reference w/ no offsets to any labels.

        :param float d_rv: small doppler shift in km/s
        :param bool symmetric: if True, applies both positive and negative doppler shifts
        :return:
        """
        warn(
            "This feature is experimental and has not been sufficiently tested on either computational or "
            "statistical grounds!",
            UserWarning,
        )
        self.labels.loc["RV"] = 0.0
        self.labels["fffff"] = self.labels["aaaaa"]
        self.labels.loc["RV", "fffff"] += d_rv
        tmp1 = doppler_shift(self.wavelength["init"], self.spectra["init"][0], d_rv)
        self.spectra["init"] = np.append(
            self.spectra["init"], tmp1[np.newaxis, :], axis=0
        )
        if symmetric:
            self.labels["ggggg"] = self.labels["aaaaa"]
            self.labels.loc["RV", "ggggg"] -= d_rv
            tmp2 = doppler_shift(
                self.wavelength["init"], self.spectra["init"][0], -d_rv
            )
            self.spectra["init"] = np.append(
                self.spectra["init"], tmp2[np.newaxis, :], axis=0
            )

    def convolve(self, instrument, name: Optional[str] = None) -> None:
        """
        Convolves spectra to instrument resolution and samples onto instrument's wavelength grid

        :param InstConfig instrument: Instrument object to convolve and sample spectra onto
        :param str name: Name to give spectra. If None, defaults to name of instrument
        :return:
        """
        if name is None:
            name = instrument.name
        outwave = instrument.wave
        self.spectra[name] = convolve_spec(
            wave=self.wavelength["init"],
            spec=self.spectra["init"],
            resolution=instrument.R_res,
            outwave=outwave,
            res_in=self.resolution["init"],
        )
        self.wavelength[name] = outwave
        self.resolution[name] = instrument.R_res

    def calc_gradient(
        self,
        name: Union[str, InstConfig],
        symmetric: bool = True,
        ref_included: bool = True,
    ) -> None:
        """
        Calculates gradients of the reference spectra with respect to each label.

        :param Union[str,InstConfig] name: Name of convolved spectra to calculate gradient for.
            Will also accept an InstConfig object and use InstConfig.name.
        :param bool symmetric: If True, calculates symmetric gradient around reference labels
        :param bool ref_included: If True, expects first spectra to be reference spectra w/ no offsets to any labels.
                                  Required for symmetric=False.
        :return:
        """
        if isinstance(name, InstConfig):
            name = name.name

        self.gradients[name] = calc_gradient(
            spectra=self.spectra[name],
            labels=self.labels,
            symmetric=symmetric,
            ref_included=ref_included,
        )
        self.gradients[name].columns = self.wavelength[name]

    def zero_gradients(
        self, name: Union[str, InstConfig], labels: Union[str, List[str]]
    ):
        """
        Sets gradients of a spectrum to zero for the specified labels. This is equivalent to setting a delta-function
        prior on those labels (i.e., holding them fixed).

        :param Union[str,InstConfig] name: Name of spectra to apply gradient zeroing to.
            Will also accept an InstConfig object and use InstConfig.name.
        :param Union[str,List[str]] labels: List of labels for which to zero gradients
        :return:
        """
        if isinstance(name, InstConfig):
            name = name.name
        self.gradients[name].loc[labels] = 0

    def mask_wavelength(
        self, name: Union[str, InstConfig], regions: List[Tuple[float, float]]
    ) -> None:
        """
        Masks the information content of a spectrum by setting the gradient to zero within the bounds of the mask.
        Can be used to mimic the masking of skylines, non-LTE lines, or detector gaps.

        :param Union[str,InstConfig] name: Name of the spectra to apply  mask to
        :param List[Tuple[float,float]] regions: List of wavelength bounds on the regions to mask.
        :return:
        """
        if isinstance(name, InstConfig):
            name = name.name
        if not isinstance(regions, list):
            regions = [regions]
        for region in regions:
            min_wave, max_wave = region
            mask = (self.wavelength[name] > min_wave) & (
                self.wavelength[name] < max_wave
            )
            self.gradients[name].iloc[:, mask] = 0

    def get_names(self) -> List[str]:
        """
        Get names of all spectra contained in this object

        :return List[str]: List of spectra names that this object contains.
        """
        return list(self.spectra.keys())

    def duplicate(self, name: str, new_name: str) -> None:
        """
        Duplicates set of spectra/gradients

        :param str name: Name of spectra to duplicate
        :param str new_name: Name given to new spectra
        :return:
        """
        self.resolution[new_name] = self.resolution[name]
        self.wavelength[new_name] = np.copy(self.wavelength[name])
        self.spectra[new_name] = np.copy(self.spectra[name])
        self.gradients[new_name] = pd.DataFrame.copy(self.gradients[name])

    def reset(self) -> None:
        """
        Resets object to only the initial high-res spectra

        :return:
        """
        init_resolution = self.resolution["init"]
        self.resolution.clear()
        self.resolution["init"] = init_resolution

        init_wavelength = self.wavelength["init"]
        self.wavelength.clear()
        self.wavelength["init"] = init_wavelength

        init_spectra = self.spectra["init"]
        self.spectra.clear()
        self.spectra["init"] = init_spectra
        self.spectra["init"] = init_spectra

        del self.gradients
        self.gradients = {}


#!/usr/bin/env python3

import sys
import math
import numpy as np
import glob
import scipy.optimize as optimization
import os
from itertools import islice
import re

#Bond Types
# 1-16 -> SP
# 17-32 -> PS
# 33-36 -> SB
# 37-38 -> BB-inter
# 39-54 -> BB-intra

WC = {
    'A': 'T',
    'C': 'G',
    'G': 'C',
    'T': 'A'
}

typeBond = {
    'SP/AA': 1,
    'SP/AC': 2,
    'SP/AG': 3,
    'SP/AT': 4,
    'SP/CA': 5,
    'SP/CC': 6,
    'SP/CG': 7,
    'SP/CT': 8,
    'SP/GA': 9,
    'SP/GC': 10,
    'SP/GG': 11,
    'SP/GT': 12,
    'SP/TA': 13,
    'SP/TC': 14,
    'SP/TG': 15,
    'SP/TT': 16,
    'PS/AA': 17,
    'PS/AC': 18,
    'PS/AG': 19,
    'PS/AT': 20,
    'PS/CA': 21,
    'PS/CC': 22,
    'PS/CG': 23,
    'PS/CT': 24,
    'PS/GA': 25,
    'PS/GC': 26,
    'PS/GG': 27,
    'PS/GT': 28,
    'PS/TA': 29,
    'PS/TC': 30,
    'PS/TG': 31,
    'PS/TT': 32,
    'SB/A': 33,
    'SB/C': 34,
    'SB/G': 35,
    'SB/T': 36,
    'BB-inter/AT': 37,
    'BB-inter/CG': 38,
    'BB-inter/GC': 38,
    'BB-inter/TA': 37,
    'BB-intra/AA': 39,
    'BB-intra/AC': 40,
    'BB-intra/AG': 41,
    'BB-intra/AT': 42,
    'BB-intra/CA': 43,
    'BB-intra/CC': 44,
    'BB-intra/CG': 45,
    'BB-intra/CT': 46,
    'BB-intra/GA': 47,
    'BB-intra/GC': 48,
    'BB-intra/GG': 49,
    'BB-intra/GT': 50,
    'BB-intra/TA': 51,
    'BB-intra/TC': 52,
    'BB-intra/TG': 53,
    'BB-intra/TT': 54
}

kbond = {
    1: 108.849600,
    2: 116.093000,
    3: 130.551200,
    4: 134.298200,
    5: 117.795600,
    6: 158.718400,
    7: 121.876200,
    8: 150.068000,
    9: 98.948800,
    10: 82.660600,
    11: 150.242600,
    12: 109.797800,
    13: 139.538000,
    14: 147.736200,
    15: 115.540000,
    16: 159.897200,
    17: 29.178200,
    18: 30.063800,
    19: 23.022000,
    20: 31.743400,
    21: 29.641600,
    22: 38.022200,
    23: 20.852400,
    24: 37.382800,
    25: 30.260400,
    26: 38.421000,
    27: 33.503600,
    28: 34.214400,
    29: 38.532200,
    30: 39.623000,
    31: 32.780400,
    32: 47.345800,
    33: 78.459400,
    34: 89.029200,
    35: 98.264400,
    36: 90.502800,
    37: 42.968000,
    38: 71.908800,
    39: 11.723300,
    40: 16.033040,
    41: 6.953480,
    42: 37.138400,
    43: 11.678180,
    44: 4.556460,
    45: 8.929340,
    46: 8.008620,
    47: 12.141360,
    48: 19.043940,
    49: 6.154040,
    50: 26.376000,
    51: 16.652460,
    52: 5.911920,
    53: 7.927460,
    54: 9.313280
}

r0bond = {
    1: 3.737680,
    2: 3.739850,
    3: 3.750370,
    4: 3.747050,
    5: 3.758460,
    6: 3.765150,
    7: 3.760470,
    8: 3.751930,
    9: 3.738060,
    10: 3.707210,
    11: 3.760770,
    12: 3.741210,
    13: 3.758390,
    14: 3.763810,
    15: 3.758300,
    16: 3.755310,
    17: 4.083580,
    18: 4.070020,
    19: 4.094660,
    20: 4.115280,
    21: 4.127790,
    22: 4.169700,
    23: 4.087590,
    24: 4.119970,
    25: 4.108820,
    26: 4.038130,
    27: 4.173860,
    28: 4.113180,
    29: 4.143020,
    30: 4.166030,
    31: 4.133930,
    32: 4.138290,
    33: 4.894800,
    34: 4.393840,
    35: 5.012620,
    36: 4.457900,
    37: 6.089870,
    38: 5.700030,
    39: 3.871730,
    40: 3.733100,
    41: 4.123320,
    42: 3.696230,
    43: 4.235810,
    44: 4.185290,
    45: 4.259420,
    46: 3.946020,
    47: 3.819480,
    48: 3.708220,
    49: 4.148530,
    50: 3.692320,
    51: 4.358370,
    52: 4.188390,
    53: 4.483310,
    54: 4.016420
}

def create_bonds(seqstring):
	lseq = len(seqstring)
	nbonds = 9*lseq-6
	bondtype = np.zeros(nbonds+1, dtype=int)
	listkbond = np.zeros(nbonds+1, dtype=float)
	listr0bond = np.zeros(nbonds+1, dtype=float)
	i1list = np.zeros(nbonds+1, dtype=int)
	i2list = np.zeros(nbonds+1, dtype=int)

	firstS1 = 1
	lastS1 = firstS1 + (lseq-1)*3
	firstP1 = 3
	lastP1 = firstP1 + (lseq-2)*3
	firstB1 = 2
	lastB1 = firstB1 + (lseq-1)*3

	firstS2 = lastB1 + 1
	lastS2 = firstS2 + (lseq-1)*3
	firstP2 = lastB1 + 3
	lastP2 = firstP2 + (lseq-2)*3
	firstB2 = lastB1 + 2
	lastB2 = firstB2 + (lseq-1)*3

	ibond = 1

	#### SP
	for i in range(1,lseq):
	    i1list[ibond] = 3*i - 2
	    i2list[ibond] = i1list[ibond] + 2
	    bondtype[ibond] = typeBond['SP/'+seqstring[i-1]+seqstring[i]]
	    listkbond[ibond] = kbond[bondtype[ibond]]
	    listr0bond[ibond] = r0bond[bondtype[ibond]]
	    ibond += 1
	for i in range(1,lseq):
	    i1list[ibond] = lastB1 + 3*i - 2
	    i2list[ibond] = i1list[ibond] + 2
	    bondtype[ibond] = typeBond['SP/'+WC[seqstring[lseq-1-(i-1)]]+WC[seqstring[lseq-1-i]]]
	    listkbond[ibond] = kbond[bondtype[ibond]]
	    listr0bond[ibond] = r0bond[bondtype[ibond]]
	    ibond += 1

	#### PS
	for i in range(2,lseq+1):
	    i1list[ibond] = 3*i - 3
	    i2list[ibond] = i1list[ibond] + 1
	    bondtype[ibond] = typeBond['PS/'+seqstring[i-2]+seqstring[i-1]]
	    listkbond[ibond] = kbond[bondtype[ibond]]
	    listr0bond[ibond] = r0bond[bondtype[ibond]]
	    ibond += 1
	for i in range(2,lseq+1):
	    i1list[ibond] = lastB1 + 3*i - 3
	    i2list[ibond] = i1list[ibond] + 1
	    bondtype[ibond] = typeBond['PS/'+WC[seqstring[lseq-1-(i-2)]]+WC[seqstring[lseq-1-(i-1)]]]
	    listkbond[ibond] = kbond[bondtype[ibond]]
	    listr0bond[ibond] = r0bond[bondtype[ibond]]
	    ibond += 1

	#### SB
	for i in range(1,lseq+1):
	    i1list[ibond] = 3*i - 2
	    i2list[ibond] = i1list[ibond] + 1
	    bondtype[ibond] = typeBond['SB/'+seqstring[i-1]]
	    listkbond[ibond] = kbond[bondtype[ibond]]
	    listr0bond[ibond] = r0bond[bondtype[ibond]]
	    ibond += 1
	for i in range(1,lseq+1):
	    i1list[ibond] = lastB1 + 3*i - 2
	    i2list[ibond] = i1list[ibond] + 1
	    bondtype[ibond] = typeBond['SB/'+WC[seqstring[lseq-1-(i-1)]]]
	    listkbond[ibond] = kbond[bondtype[ibond]]
	    listr0bond[ibond] = r0bond[bondtype[ibond]]
	    ibond += 1

	#### BB-inter
	for i in range(1,lseq+1):
	    i1list[ibond] = 3*i - 1
	    i2list[ibond] = lastB2 + 2 - i1list[ibond]
	    bondtype[ibond] = typeBond['BB-inter/'+seqstring[i-1]+WC[seqstring[i-1]]]
	    listkbond[ibond] = kbond[bondtype[ibond]]
	    listr0bond[ibond] = r0bond[bondtype[ibond]]
	    ibond += 1

	#### BB-intra
	for i in range(1,lseq):
	    i1list[ibond] = 3*i - 1
	    i2list[ibond] = i1list[ibond] + 3
	    bondtype[ibond] = typeBond['BB-intra/'+seqstring[i-1]+seqstring[i]]
	    listkbond[ibond] = kbond[bondtype[ibond]]
	    listr0bond[ibond] = r0bond[bondtype[ibond]]
	    ibond += 1
	for i in range(1,lseq):
	    i1list[ibond] = lastB1 + 3*i - 1
	    i2list[ibond] = i1list[ibond] + 3
	    bondtype[ibond] = typeBond['BB-intra/'+WC[seqstring[lseq-1-(i-1)]]+WC[seqstring[lseq-1-i]]]
	    listkbond[ibond] = kbond[bondtype[ibond]]
	    listr0bond[ibond] = r0bond[bondtype[ibond]]
	    ibond += 1

	return i1list, i2list, listr0bond, listkbond