# coding: utf-8
# Copyright (c) Pymatgen Development Team.
# Distributed under the terms of the MIT License.

from __future__ import division, unicode_literals

import six
import ruamel.yaml as yaml
import os
import json

"""
This module provides classes to perform analyses of
the local environments (e.g., finding near neighbors)
of single sites in molecules and structures.
To do:
- Insert LocalStructOrderParas class here.
"""

__author__ = "Shyue Ping Ong, Geoffroy Hautier, Sai Jayaraman,"+\
    " Nils E. R. Zimmermann, Bharat Medasani"
__copyright__ = "Copyright 2011, The Materials Project"
__version__ = "1.0"
__maintainer__ = "Nils E. R. Zimmermann"
__email__ = "nils.e.r.zimmermann@gmail.com"
__status__ = "Production"
__date__ = "August 17, 2017"

from math import pow, pi, asin, atan, sqrt, exp, cos, acos
import numpy as np

from bisect import bisect_left
from scipy.spatial import Voronoi
from pymatgen import Element
from pymatgen.core.structure import Structure
from pymatgen.util.num import abs_cap
from pymatgen.analysis.bond_valence import BV_PARAMS
from pymatgen.analysis.structure_analyzer import OrderParameters

file_dir = os.path.dirname(__file__)
rad_file = os.path.join(file_dir, 'ionic_radii.json')
with open(rad_file, 'r') as fp:
    _ion_radii = json.load(fp)

class ValenceIonicRadiusEvaluator(object):
    """
    Computes site valences and ionic radii for a structure using bond valence
    analyzer

    Args:
        structure: pymatgen.core.structure.Structure
    """

    def __init__(self, structure):
        self._structure = structure.copy()
        self._valences = self._get_valences()
        self._ionic_radii = self._get_ionic_radii()

    @property
    def radii(self):
        """
        List of ionic radii of elements in the order of sites.
        """
        el = [site.species_string for site in self._structure.sites]
        radii_dict = dict(zip(el, self._ionic_radii))
        #print radii_dict
        return radii_dict

    @property
    def valences(self):
        """
        List of oxidation states of elements in the order of sites.
        """
        el = [site.species_string for site in self._structure.sites]
        valence_dict = dict(zip(el, self._valences))
        return valence_dict

    @property
    def structure(self):
        """
        Returns oxidation state decorated structure.
        """
        return self._structure.copy()

    def _get_ionic_radii(self):
        """
        Computes ionic radii of elements for all sites in the structure.
        If valence is zero, atomic radius is used.
        """
        radii = []
        vnn = VoronoiNN() # self._structure)

        def nearest_key(sorted_vals, key):
            i = bisect_left(sorted_vals, key)
            if i == len(sorted_vals):
                return sorted_vals[-1]
            if i == 0:
                return sorted_vals[0]
            before = sorted_vals[i-1]
            after = sorted_vals[i]
            if after-key < key-before:
                return after
            else:
                return before

        for i in range(len(self._structure.sites)):
            site = self._structure.sites[i]
            if isinstance(site.specie,Element):
                radius = site.specie.atomic_radius
                # Handle elements with no atomic_radius
                # by using calculated values instead.
                if radius is None:
                    radius = site.specie.atomic_radius_calculated
                if radius is None:
                    raise ValueError(
                            "cannot assign radius to element {}".format(
                            site.specie))
                radii.append(radius)
                continue

            el = site.specie.symbol
            oxi_state = int(round(site.specie.oxi_state))
            coord_no = int(round(vnn.get_cn(self._structure, i)))
            try:
                tab_oxi_states = sorted(map(int, _ion_radii[el].keys()))
                oxi_state = nearest_key(tab_oxi_states, oxi_state)
                radius = _ion_radii[el][str(oxi_state)][str(coord_no)]
            except KeyError:
                if vnn.get_cn(self._structure, i)-coord_no > 0:
                    new_coord_no = coord_no + 1
                else:
                    new_coord_no = coord_no - 1
                try:
                    radius = _ion_radii[el][str(oxi_state)][str(new_coord_no)]
                    coord_no = new_coord_no
                except:
                    tab_coords = sorted(map(int, _ion_radii[el][str(oxi_state)].keys()))
                    new_coord_no = nearest_key(tab_coords, coord_no)
                    i = 0
                    for val in tab_coords:
                        if  val > coord_no:
                            break
                        i = i + 1
                    if i == len(tab_coords):
                        key = str(tab_coords[-1])
                        radius = _ion_radii[el][str(oxi_state)][key]
                    elif i == 0:
                        key = str(tab_coords[0])
                        radius = _ion_radii[el][str(oxi_state)][key]
                    else:
                        key = str(tab_coords[i-1])
                        radius1 = _ion_radii[el][str(oxi_state)][key]
                        key = str(tab_coords[i])
                        radius2 = _ion_radii[el][str(oxi_state)][key]
                        radius = (radius1+radius2)/2

            #implement complex checks later
            radii.append(radius)
        return radii

    def _get_valences(self):
        """
        Computes ionic valences of elements for all sites in the structure.
        """
        try:
            bv = BVAnalyzer()
            self._structure = bv.get_oxi_state_decorated_structure(self._structure)
            valences = bv.get_valences(self._structure)
        except:
            try:
                bv = BVAnalyzer(symm_tol=0.0)
                self._structure = bv.get_oxi_state_decorated_structure(self._structure)
                valences = bv.get_valences(self._structure)
            except:
                valences = []
                for site in self._structure.sites:
                    if len(site.specie.common_oxidation_states) > 0:
                        valences.append(site.specie.common_oxidation_states[0])
                    # Handle noble gas species
                    # which have no entries in common_oxidation_states.
                    else:
                        valences.append(0)
                if sum(valences):
                    valences = [0]*self._structure.num_sites
                else:
                    self._structure.add_oxidation_state_by_site(valences)
                #raise

        #el = [site.specie.symbol for site in self._structure.sites]
        #el = [site.species_string for site in self._structure.sites]
        #el = [site.specie for site in self._structure.sites]
        #valence_dict = dict(zip(el, valences))
        #print valence_dict
        return valences

class NearNeighbors(object):
    """
    Base class to determine near neighbors that typically include nearest
    neighbors and others that are within some tolerable distance.
    """

    def __init__(self):
        pass

    def get_cn(self, structure, n, use_weights=False):
        """
        Get coordination number, CN, of site with index n in structure.

        Args:
            structure (Structure): input structure.
            n (integer): index of site for which to determine CN.
            use_weights (boolean): flag indicating whether (True)
                to use weights for computing the coordination number
                or not (False, default: each coordinated site has equal
                weight).
        Returns:
            cn (integer or float): coordination number.
        """

        siw = self.get_nn_info(structure, n)
        return sum([e['weight'] for e in siw]) if use_weights else len(siw)

    def get_nn(self, structure, n):
        """
        Get near neighbors of site with index n in structure.

        Args:
            structure (Structure): input structure.
            n (integer): index of site in structure for which to determine
                    neighbors.
        Returns:
            sites (list of Site objects): near neighbors.
        """

        return [e['site'] for e in self.get_nn_info(structure, n)]

    def get_weights_of_nn_sites(self, n):
        """
        Get weight associated with each near neighbor of site with
        index n in structure.

        Args:
            structure (Structure): input structure.
            n (integer): index of site for which to determine the weights.
        Returns:
            weights (list of floats): near-neighbor weights.
        """

        return [e['weight'] for e in self.get_nn_info(structure, n)]

    def get_nn_images(self, structure, n):
        """
        Get image location of all near neighbors of site with index n in
        structure.

        Args:
            structure (Structure): input structure.
            n (integer): index of site for which to determine the image
                location of near neighbors.
        Returns:
            images (list of 3D integer array): image locations of
                near neighbors.
        """

        return [e['image'] for e in self.get_nn_info(structure, n)]

    def get_nn_info(self, structure, n):
        """
        Get all near-neighbor sites as well as the associated image locations
        and weights of the site with index n.

        Args:
            structure (Structure): input structure.
            n (integer): index of site for which to determine near-neighbor
                information.

        Returns:
            siw (list of dicts): each dictionary provides information
                about a single near neighbor, where key 'site' gives
                access to the corresponding Site object, 'image' gives
                the image location, and 'weight' provides the weight
                that a given near-neighbor site contributes
                to the coordination number (1 or smaller), 'site_index'
                gives index of the corresponding site in
                the original structure.
        """

        raise NotImplementedError("get_nn_info(structure, n)"
                " is not defined!")

    @staticmethod
    def _get_image(frac_coords):
        """Private convenience method for get_nn_info,
        gives lattice image from provided PeriodicSite."""
        return [int(f) if f >= 0 else int(f - 1)
                for f in frac_coords]

    @staticmethod
    def _get_original_site(structure, site):
        """Private convenience method for get_nn_info,
        gives original site index from ProvidedPeriodicSite."""
        is_periodic_image = [site.is_periodic_image(s) for s in structure]
        return is_periodic_image.index(True)

class VoronoiNN(NearNeighbors):
    """
    Uses a Voronoi algorithm to determine near neighbors for each site in a
    structure.

    Args:
        tol (float): tolerance parameter for near-neighbor finding
            (default: 0).
        targets (Element or list of Elements): target element(s).
        cutoff (float): cutoff radius in Angstrom to look for near-neighbor
            atoms. Defaults to 10.0.
        allow_pathological (bool): whether to allow infinite vertices in
            determination of Voronoi coordination.
    """

    def __init__(self, tol=0, targets=None, cutoff=10.0,
                 allow_pathological=False):
        self.tol = tol
        self.cutoff = cutoff
        self.allow_pathological = allow_pathological
        self.targets = targets

    def get_voronoi_polyhedra(self, structure, n):
        """
        Gives a weighted polyhedra around a site. This uses the Voronoi
        construction with solid angle weights.
        See ref: A Proposed Rigorous Definition of Coordination Number,
        M. O'Keeffe, Acta Cryst. (1979). A35, 772-775

        Args:
            structure (Structure): structure for which to evaluate the
                coordination environment.
            n (integer): site index.

        Returns:
            A dict of sites sharing a common Voronoi facet with the site
            n and their solid angle weights
        """
        if self.targets is None:
            targets = structure.composition.elements
        else:
            targets = self.targets
        center = structure[n]
        neighbors = structure.get_sites_in_sphere(
            center.coords, self.cutoff)
        neighbors = [i[0] for i in sorted(neighbors, key=lambda s: s[1])]
        qvoronoi_input = [s.coords for s in neighbors]
        voro = Voronoi(qvoronoi_input)
        all_vertices = voro.vertices

        results = {}
        for nn, vind in voro.ridge_dict.items():
            if 0 in nn:
                if -1 in vind:
                    if self.allow_pathological:
                        continue
                    else:
                        raise RuntimeError("This structure is pathological,"
                                           " infinite vertex in the voronoi "
                                           "construction")

                facets = [all_vertices[i] for i in vind]
                results[neighbors[sorted(nn)[1]]] = solid_angle(
                    center.coords, facets)

        maxangle = max(results.values())

        resultweighted = {}
        for nn, angle in results.items():
            # is nn site is ordered use "nn.specie" to get species, else use "nn.species_and_occu" to get species
            if nn.is_ordered:
                if nn.specie in targets:
                    resultweighted[nn] = angle / maxangle
            else:  # is nn site is disordered
                for disordered_sp in nn.species_and_occu.keys():
                    if disordered_sp in targets:
                        resultweighted[nn] = angle / maxangle

        return resultweighted

    def get_nn_info(self, structure, n):
        """"
        Get all near-neighbor sites as well as the associated image locations
        and weights of the site with index n in structure
        using Voronoi decomposition.

        Args:
            structure (Structure): input structure.
            n (integer): index of site for which to determine near-neighbor
                sites.

        Returns:
            siw (list of tuples (Site, array, float)): tuples, each one
                of which represents a coordinated site, its image location,
                and its weight.
        """

        if self.targets is None:
            targets = structure.composition.elements
        else:
            targets = self.targets
        siw = []
        for site, weight in self.get_voronoi_polyhedra(
                structure, n).items():
            if weight > self.tol and site.specie in targets:
                siw.append({'site': site,
                            'image': self._get_image(site.frac_coords),
                            'weight': weight,
                            'site_index': self._get_original_site(structure, site)})
        return siw

class JMolNN(NearNeighbors):
    """
    Determine near-neighbor sites and coordination number using an emulation
    of JMol's default autoBond() algorithm. This version of the algorithm
    does not take into account any information regarding known charge
    states.

    Args:
        tol (float): tolerance parameter for bond determination
            (default: 1E-3).
        el_radius_updates: (dict) symbol->float to override default atomic
            radii table values
    """

    def __init__(self, tol=1E-3, el_radius_updates=None):

        self.tol = tol

        # Load elemental radii table
        bonds_file = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                                  "bonds_jmol_ob.yaml")
        with open(bonds_file, 'r') as f:
            self.el_radius = yaml.safe_load(f)

        # Update any user preference elemental radii
        if el_radius_updates:
            self.el_radius.update(el_radius_updates)

    def get_max_bond_distance(self, el1_sym, el2_sym, constant=0.56):
        """
        Use JMol algorithm to determine bond length from atomic parameters
        Args:
            el1_sym: (str) symbol of atom 1
            el2_sym: (str) symbol of atom 2
            constant: (float) factor to tune model

        Returns: (float) max bond length

        """
        return sqrt(
            (self.el_radius[el1_sym] + self.el_radius[el2_sym] + constant) ** 2)

    def get_nn_info(self, structure, n):
        """
        Get all near-neighbor sites as well as the associated image locations
        and weights of the site with index n using the bond identification
        algorithm underlying JMol.

        Args:
            structure (Structure): input structure.
            n (integer): index of site for which to determine near
                neighbors.

        Returns:
            siw (list of tuples (Site, array, float)): tuples, each one
                of which represents a neighbor site, its image location,
                and its weight.
        """

        site = structure[n]

        # Determine relevant bond lengths based on atomic radii table
        bonds = {}
        for el in structure.composition.elements:
            bonds[site.specie, el] = self.get_max_bond_distance(
                site.specie.symbol, el.symbol)

        # Search for neighbors up to max bond length + tolerance
        max_rad = max(bonds.values()) + self.tol
        min_rad = min(bonds.values())

        siw = []
        for neighb, dist in structure.get_neighbors(site, max_rad):
            # Confirm neighbor based on bond length specific to atom pair
            if dist <= bonds[(site.specie, neighb.specie)] + self.tol:
                weight = min_rad / dist
                siw.append({'site': neighb,
                            'image': self._get_image(neighb.frac_coords),
                            'weight': weight,
                            'site_index': self._get_original_site(structure, neighb)})
        return siw

class MinimumDistanceNN(NearNeighbors):
    """
    Determine near-neighbor sites and coordination number using the
    nearest neighbor(s) at distance, d_min, plus all neighbors
    within a distance (1 + delta) * d_min, where delta is a
    (relative) distance tolerance parameter.

    Args:
        tol (float): tolerance parameter for neighbor identification
            (default: 0.1).
        cutoff (float): cutoff radius in Angstrom to look for trial
            near-neighbor sites (default: 10.0).
    """

    def __init__(self, tol=0.1, cutoff=10.0):

        self.tol = tol
        self.cutoff = cutoff

    def get_nn_info(self, structure, n):
        """
        Get all near-neighbor sites as well as the associated image locations
        and weights of the site with index n using the closest neighbor
        distance-based method.

        Args:
            structure (Structure): input structure.
            n (integer): index of site for which to determine near
                neighbors.

        Returns:
            siw (list of tuples (Site, array, float)): tuples, each one
                of which represents a neighbor site, its image location,
                and its weight.
        """

        site = structure[n]
        neighs_dists = structure.get_neighbors(site, self.cutoff)
        min_dist = min([dist for neigh, dist in neighs_dists])

        siw = []
        for s, dist in neighs_dists:
            if dist < (1.0 + self.tol) * min_dist:
                w = min_dist / dist
                siw.append({'site': s,
                            'image': self._get_image(s.frac_coords),
                            'weight': w,
                            'site_index': self._get_original_site(structure, s)})
        return siw

class MinimumOKeeffeNN(NearNeighbors):
    """
    Determine near-neighbor sites and coordination number using the
    neighbor(s) at closest relative distance, d_min_OKeffee, plus some
    relative tolerance, where bond valence parameters from O'Keeffe's
    bond valence method (J. Am. Chem. Soc. 1991, 3226-3229) are used
    to calculate relative distances.

    Args:
        tol (float): tolerance parameter for neighbor identification
            (default: 0.1).
        cutoff (float): cutoff radius in Angstrom to look for trial
            near-neighbor sites (default: 10.0).
    """

    def __init__(self, tol=0.1, cutoff=10.0):

        self.tol = tol
        self.cutoff = cutoff

    def get_nn_info(self, structure, n):
        """
        Get all near-neighbor sites as well as the associated image locations
        and weights of the site with index n using the closest relative
        neighbor distance-based method with O'Keeffe parameters.

        Args:
            structure (Structure): input structure.
            n (integer): index of site for which to determine near
                neighbors.

        Returns:
            siw (list of tuples (Site, array, float)): tuples, each one
                of which represents a neighbor site, its image location,
                and its weight.
        """

        site = structure[n]
        neighs_dists = structure.get_neighbors(site, self.cutoff)
        try:
            eln = site.specie.element
        except:
            eln = site.species_string

        reldists_neighs = []
        for neigh, dist in neighs_dists:
            try:
                el2 = neigh.specie.element
            except:
                el2 = neigh.species_string
            reldists_neighs.append([dist / get_okeeffe_distance_prediction(
                    eln, el2), neigh])

        siw = []
        min_reldist = min([reldist for reldist, neigh in reldists_neighs])
        for reldist, s in reldists_neighs:
            if reldist < (1.0 + self.tol) * min_reldist:
                w = min_reldist / reldist
                siw.append({'site': s,
                            'image': self._get_image(s.frac_coords),
                            'weight': w,
                            'site_index': self._get_original_site(structure, s)})

        return siw

class MinimumVIRENN(NearNeighbors):
    """
    Determine near-neighbor sites and coordination number using the
    neighbor(s) at closest relative distance, d_min_VIRE, plus some
    relative tolerance, where atom radii from the
    ValenceIonicRadiusEvaluator (VIRE) are used
    to calculate relative distances.

    Args:
        tol (float): tolerance parameter for neighbor identification
            (default: 0.1).
        cutoff (float): cutoff radius in Angstrom to look for trial
            near-neighbor sites (default: 10.0).
    """

    def __init__(self, tol=0.1, cutoff=10.0):

        self.tol = tol
        self.cutoff = cutoff

    def get_nn_info(self, structure, n):
        """
        Get all near-neighbor sites as well as the associated image locations
        and weights of the site with index n using the closest relative
        neighbor distance-based method with VIRE atomic/ionic radii.

        Args:
            structure (Structure): input structure.
            n (integer): index of site for which to determine near
                neighbors.

        Returns:
            siw (list of tuples (Site, array, float)): tuples, each one
                of which represents a neighbor site, its image location,
                and its weight.
        """

        vire = ValenceIonicRadiusEvaluator(structure)
        site = vire.structure[n]
        neighs_dists = vire.structure.get_neighbors(site, self.cutoff)
        rn = vire.radii[vire.structure[n].species_string]

        reldists_neighs = []
        for neigh, dist in neighs_dists:
            reldists_neighs.append([dist / (
                    vire.radii[neigh.species_string] + rn), neigh])

        siw = []
        min_reldist = min([reldist for reldist, neigh in reldists_neighs])
        for reldist, s in reldists_neighs:
            if reldist < (1.0 + self.tol) * min_reldist:
                w = min_reldist / reldist
                siw.append({'site': s,
                            'image': self._get_image(s.frac_coords),
                            'weight': w,
                            'site_index': self._get_original_site(structure, s)})

        return siw

def solid_angle(center, coords):
    """
    Helper method to calculate the solid angle of a set of coords from the
    center.

    Args:
        center (3x1 array): Center to measure solid angle from.
        coords (Nx3 array): List of coords to determine solid angle.

    Returns:
        The solid angle.
    """
    o = np.array(center)
    r = [np.array(c) - o for c in coords]
    r.append(r[0])
    n = [np.cross(r[i + 1], r[i]) for i in range(len(r) - 1)]
    n.append(np.cross(r[1], r[0]))
    vals = []
    for i in range(len(n) - 1):
        v = -np.dot(n[i], n[i + 1]) \
            / (np.linalg.norm(n[i]) * np.linalg.norm(n[i + 1]))
        vals.append(acos(abs_cap(v)))
    phi = sum(vals)
    return phi + (3 - len(r)) * pi

def get_okeeffe_params(el_symbol):
    """
    Returns the elemental parameters related to atom size and
    electronegativity which are used for estimating bond-valence
    parameters (bond length) of pairs of atoms on the basis of data
    provided in 'Atoms Sizes and Bond Lengths in Molecules and Crystals'
    (O'Keeffe & Brese, 1991).

    Args:
        el_symbol (str): element symbol.
    Returns:
        (dict): atom-size ('r') and electronegativity-related ('c')
                parameter.
    """

    el = Element(el_symbol)
    if el not in list(BV_PARAMS.keys()):
        raise RuntimeError("Could not find O'Keeffe parameters for element"
                           " \"{}\" in \"BV_PARAMS\"dictonary"
                           " provided by pymatgen".format(el_symbol))

    return BV_PARAMS[el]

def get_okeeffe_distance_prediction(el1, el2):
    """
    Returns an estimate of the bond valence parameter (bond length) using
    the derived parameters from 'Atoms Sizes and Bond Lengths in Molecules
    and Crystals' (O'Keeffe & Brese, 1991). The estimate is based on two
    experimental parameters: r and c. The value for r  is based off radius,
    while c is (usually) the Allred-Rochow electronegativity. Values used
    are *not* generated from pymatgen, and are found in
    'okeeffe_params.json'.

    Args:
        el1, el2 (Element): two Element objects
    Returns:
        a float value of the predicted bond length
    """
    el1_okeeffe_params = get_okeeffe_params(el1)
    el2_okeeffe_params = get_okeeffe_params(el2)

    r1 = el1_okeeffe_params['r']
    r2 = el2_okeeffe_params['r']
    c1 = el1_okeeffe_params['c']
    c2 = el2_okeeffe_params['c']

    return r1 + r2 - r1 * r2 * pow(
            sqrt(c1) - sqrt(c2), 2) / (c1 * r1 + c2 * r2)

def get_neighbors_of_site_with_index(struct, n, approach="min_dist", delta=0.1, \
        cutoff=10.0):
    """
    Returns the neighbors of a given site using a specific neighbor-finding
    method.

    Args:
        struct (Structure): input structure.
        n (int): index of site in Structure object for which motif type
                is to be determined.
        approach (str): type of neighbor-finding approach, where
              "min_dist" will use the MinimumDistanceNN class,
              "voronoi" the VoronoiNN class, "min_OKeeffe" the
              MinimumOKeeffe class, and "min_VIRE" the MinimumVIRENN class.
        delta (float): tolerance involved in neighbor finding.
        cutoff (float): (large) radius to find tentative neighbors.

    Returns: neighbor sites.
    """

    if approach == "min_dist":
        return MinimumDistanceNN(tol=delta, cutoff=cutoff).get_nn(
                struct, n)
    elif approach == "voronoi":
        return VoronoiNN(tol=delta, cutoff=cutoff).get_nn(
                struct, n)
    elif approach == "min_OKeeffe":
        return MinimumOKeeffeNN(tol=delta, cutoff=cutoff).get_nn(
                struct, n)
    elif approach == "min_VIRE":
        return MinimumVIRENN(tol=delta, cutoff=cutoff).get_nn(
                struct, n)
    else:
        raise RuntimeError("unsupported neighbor-finding method ({}).".format(
                approach))

def site_is_of_motif_type(struct, n, approach="min_dist", delta=0.1, \
        cutoff=10.0, thresh=None):
    """
    Returns the motif type of the site with index n in structure struct;
    currently featuring "tetrahedral", "octahedral", "bcc", and "cp"
    (close-packed: fcc and hcp) as well as "square pyramidal" and
    "trigonal bipyramidal".  If the site is not recognized,
    "unrecognized" is returned.  If a site should be assigned to two
    different motifs, "multiple assignments" is returned.

    Args:
        struct (Structure): input structure.
        n (int): index of site in Structure object for which motif type
                is to be determined.
        approach (str): type of neighbor-finding approach, where
              "min_dist" will use the MinimumDistanceNN class,
              "voronoi" the VoronoiNN class, "min_OKeeffe" the
              MinimumOKeeffe class, and "min_VIRE" the MinimumVIRENN class.
        delta (float): tolerance involved in neighbor finding.
        cutoff (float): (large) radius to find tentative neighbors.
        thresh (dict): thresholds for motif criteria (currently, required
                keys and their default values are "qtet": 0.5,
                "qoct": 0.5, "qbcc": 0.5, "q6": 0.4).

    Returns: motif type (str).
    """

    if thresh is None:
        thresh = {
            "qtet": 0.5, "qoct": 0.5, "qbcc": 0.5, "q6": 0.4,
            "qtribipyr": 0.8, "qsqpyr": 0.8}

    ops = OrderParameters([
            "cn", "tet", "oct", "bcc", "q6", "sq_pyr", "tri_bipyr"])

    neighs_cent = get_neighbors_of_site_with_index(
            struct, n, approach=approach, delta=delta, cutoff=cutoff)
    neighs_cent.append(struct.sites[n])
    opvals = ops.get_order_parameters(
            neighs_cent, len(neighs_cent)-1, indices_neighs=[
            i for i in range(len(neighs_cent)-1)])
    cn = int(opvals[0] + 0.5)
    motif_type = "unrecognized"
    nmotif = 0

    if cn == 4 and opvals[1] > thresh["qtet"]:
        motif_type = "tetrahedral"
        nmotif += 1
    if cn == 5 and opvals[5] > thresh["qsqpyr"]:
       motif_type = "square pyramidal"
       nmotif += 1
    if cn == 5 and opvals[6] > thresh["qtribipyr"]:
       motif_type = "trigonal bipyramidal"
       nmotif += 1
    if cn == 6 and opvals[2] > thresh["qoct"]:
        motif_type = "octahedral"
        nmotif += 1
    if cn == 8 and (opvals[3] > thresh["qbcc"] and opvals[1] < thresh["qtet"]):
        motif_type = "bcc"
        nmotif += 1
    if cn == 12 and (opvals[4] > thresh["q6"] and opvals[1] < thresh["q6"] and \
                                 opvals[2] < thresh["q6"] and opvals[3] < thresh["q6"]):
        motif_type = "cp"
        nmotif += 1

    if nmotif > 1:
        motif_type = "multiple assignments"

    return motif_type


import os
import subprocess
import numpy as np
import matplotlib.pyplot as pyplot
import matplotlib.cm as cm
from matplotlib.colors import Normalize
from matplotlib.backends.backend_pdf import PdfPages
from mpl_toolkits.axes_grid1 import make_axes_locatable
from simtk import unit
import openmmtools
from cg_openmm.utilities.util import set_box_vectors, get_box_vectors
from simtk.openmm.app.pdbfile import PDBFile
from simtk.openmm.app.dcdfile import DCDFile
from mdtraj.formats import PDBTrajectoryFile
from mdtraj import Topology, Trajectory
from pymbar import timeseries
from scipy.special import erf
from scipy.optimize import minimize_scalar
import time

from openmmtools.multistate import MultiStateReporter, MultiStateSampler, ReplicaExchangeSampler
from openmmtools.multistate import ReplicaExchangeAnalyzer

# quiet down some citation spam
MultiStateSampler._global_citation_silence = True

kB = (unit.MOLAR_GAS_CONSTANT_R).in_units_of(unit.kilojoule / (unit.kelvin * unit.mole))

def make_replica_dcd_files(
    topology, timestep=5*unit.femtosecond, time_interval=200,
    output_dir="output", output_data="output.nc", checkpoint_data="output_checkpoint.nc",
    frame_begin=0, frame_stride=1):
    """
    Make dcd files from replica exchange simulation trajectory data.

    :param topology: OpenMM Topology
    :type topology: `Topology() <https://simtk.org/api_docs/openmm/api4_1/python/classsimtk_1_1openmm_1_1app_1_1topology_1_1Topology.html>`_

    :param timestep: Time step used in the simulation (default=5*unit.femtosecond)
    :type timestep: `Quantity() <http://docs.openmm.org/development/api-python/generated/simtk.unit.quantity.Quantity.html>` float * simtk.unit

    :param time_interval: frequency, in number of time steps, at which positions were recorded (default=200)
    :type time_interval: int

    :param output_dir: path to which we will write the output (default='output')
    :type output_dir: str

    :param output_data: name of output .nc data file (default='output.nc')
    :type output_data: str

    :param checkpoint_data: name of checkpoint .nc data file (default='output_checkpoint.nc')
    :type checkpoint_data: str

    :param frame_begin: Frame at which to start writing the dcd trajectory (default=0)
    :type frame_begin: int

    :param frame_stride: advance by this many time intervals when writing dcd trajectories (default=1)
    :type frame_stride: int
    """

    file_list = []

    output_data_path = os.path.join(output_dir, output_data)

    # Get number of replicas:
    reporter = MultiStateReporter(output_data_path, open_mode="r")
    states = reporter.read_thermodynamic_states()[0]
    n_replicas=len(states)

    sampler_states = reporter.read_sampler_states(iteration=0)
    xunit = sampler_states[0].positions[0].unit

    for replica_index in range(n_replicas):
        replica_positions = extract_trajectory(topology, replica_index=replica_index,
            output_data=output_data_path, checkpoint_data=checkpoint_data,
            frame_begin=frame_begin, frame_stride=frame_stride)

        n_frames_tot = replica_positions.shape[0]

        # Determine simulation time (in ps) for each frame:
        time_delta_ps = (timestep*time_interval).value_in_unit(unit.picosecond)
        traj_times = np.linspace(
            frame_begin*time_delta_ps,
            (frame_begin+frame_stride*(n_frames_tot-1))*time_delta_ps,
            num=n_frames_tot,
        )

        file_name = f"{output_dir}/replica_{replica_index+1}.dcd"

        # Trajectories are written in nanometers:
        replica_traj = Trajectory(
            replica_positions,
            Topology.from_openmm(topology),
            time=traj_times,
        )

        Trajectory.save_dcd(replica_traj,file_name)

    return file_list

def make_replica_pdb_files(
    topology, output_dir="output", output_data="output.nc", checkpoint_data="output_checkpoint.nc",
    frame_begin=0, frame_stride=1):
    """
    Make pdb files from replica exchange simulation trajectory data.

    :param topology: OpenMM Topology
    :type topology: `Topology() <https://simtk.org/api_docs/openmm/api4_1/python/classsimtk_1_1openmm_1_1app_1_1topology_1_1Topology.html>`_

    :param output_dir: path to which we will write the output (default='output')
    :type output_dir: str

    :param output_data: name of output .nc data file (default='output.nc')
    :type output_data: str

    :param checkpoint_data: name of checkpoint .nc data file (default='output_checkpoint.nc')
    :type checkpoint_data: str

    :param frame_begin: Frame at which to start writing the pdb trajectory (default=0)
    :type frame_begin: int

    :param frame_stride: advance by this many frames when writing pdb trajectories (default=1)
    :type frame_stride: int

    :returns:
        - file_list ( List( str ) ) - A list of names for the files that were written
    """
    file_list = []

    output_data_path = os.path.join(output_dir, output_data)

    # Get number of replicas:
    reporter = MultiStateReporter(output_data_path, open_mode="r")
    states = reporter.read_thermodynamic_states()[0]
    n_replicas = len(states)

    sampler_states = reporter.read_sampler_states(iteration=0)
    xunit = sampler_states[0].positions[0].unit

    for replica_index in range(n_replicas):
        replica_positions = extract_trajectory(topology, replica_index=replica_index,
            output_data=output_data_path, checkpoint_data=checkpoint_data,
            frame_begin=frame_begin, frame_stride=frame_stride)

        file_name = f"{output_dir}/replica_{replica_index+1}.pdb"

        # Trajectories are written in nanometers:
        replica_traj = Trajectory(
            replica_positions,
            Topology.from_openmm(topology),
        )

        Trajectory.save_pdb(replica_traj,file_name)

    return file_list

def make_state_dcd_files(
    topology, timestep=5*unit.femtosecond, time_interval=200,
    output_dir="output", output_data="output.nc", checkpoint_data="output_checkpoint.nc",
    frame_begin=0, frame_stride=1, center=True):
    """
    Make dcd files by state from replica exchange simulation trajectory data.
    Note: these are discontinuous trajectories with constant temperature state.

    :param topology: OpenMM Topology
    :type topology: `Topology() <https://simtk.org/api_docs/openmm/api4_1/python/classsimtk_1_1openmm_1_1app_1_1topology_1_1Topology.html>`_

    :param timestep: Time step used in the simulation (default=5*unit.femtosecond)
    :type timestep: `Quantity() <http://docs.openmm.org/development/api-python/generated/simtk.unit.quantity.Quantity.html>` float * simtk.unit

    :param time_interval: frequency, in number of time steps, at which positions were recorded (default=200)
    :type time_interval: int

    :param output_dir: path to which we will write the output (default='output')
    :type output_dir: str

    :param output_data: name of output .nc data file (default='output.nc')
    :type output_data: str

    :param checkpoint_data: name of checkpoint .nc data file (default='output_checkpoint.nc')
    :type checkpoint_data: str

    :param frame_begin: Frame at which to start writing the dcd trajectory (default=0)
    :type frame_begin: int

    :param frame_stride: advance by this many time intervals when writing dcd trajectories (default=1)
    :type frame_stride: int

    :param center: align the center of mass of each structure in the discontinuous state trajectory (default=True)
    :type center: Boolean

    """

    file_list = []

    output_data_path = os.path.join(output_dir, output_data)

    # Get number of states:
    reporter = MultiStateReporter(output_data_path, open_mode="r")
    states = reporter.read_thermodynamic_states()[0]

    sampler_states = reporter.read_sampler_states(iteration=0)
    xunit = sampler_states[0].positions[0].unit

    for state_index in range(len(states)):
        state_positions = extract_trajectory(topology, state_index=state_index,
            output_data=output_data_path, checkpoint_data=checkpoint_data,
            frame_begin=frame_begin, frame_stride=frame_stride)

        n_frames_tot = state_positions.shape[0]

        # Determine simulation time (in ps) for each frame:
        time_delta_ps = (timestep*time_interval).value_in_unit(unit.picosecond)
        traj_times = np.linspace(
            frame_begin*time_delta_ps,
            (frame_begin+frame_stride*(n_frames_tot-1))*time_delta_ps,
            num=n_frames_tot,
        )

        file_name = f"{output_dir}/state_{state_index+1}.dcd"

        # Trajectories are written in nanometers:
        state_traj = Trajectory(
            state_positions,
            Topology.from_openmm(topology),
            time=traj_times,
        )

        if center:
            ref_traj = state_traj[0]
            state_traj.superpose(ref_traj)
            # This rewrites to state_traj

        Trajectory.save_dcd(state_traj,file_name)

    return file_list

def make_state_pdb_files(
    topology, output_dir="output", output_data="output.nc", checkpoint_data="output_checkpoint.nc",
    frame_begin=0, frame_stride=1, center=True):
    """
    Make pdb files by state from replica exchange simulation trajectory data.
    Note: these are discontinuous trajectories with constant temperature state.

    :param topology: OpenMM Topology
    :type topology: `Topology() <https://simtk.org/api_docs/openmm/api4_1/python/classsimtk_1_1openmm_1_1app_1_1topology_1_1Topology.html>`_

    :param output_dir: path to which we will write the output (default='output')
    :type output_dir: str

    :param output_data: name of output .nc data file (default='output.nc')
    :type output_data: str

    :param checkpoint_data: name of checkpoint .nc data file (default='output_checkpoint.nc')
    :type checkpoint_data: str

    :param frame_begin: Frame at which to start writing the pdb trajectory (default=0)
    :type frame_begin: int

    :param frame_stride: advance by this many frames when writing pdb trajectories (default=1)
    :type frame_stride: int

    :param center: align the center of mass of each structure in the discontinuous state trajectory (default=True)
    :type center: Boolean

    :returns:
        - file_list ( List( str ) ) - A list of names for the files that were written
    """
    file_list = []

    output_data_path = os.path.join(output_dir, output_data)

    # Get number of states:
    reporter = MultiStateReporter(output_data_path, open_mode="r")
    states = reporter.read_thermodynamic_states()[0]

    sampler_states = reporter.read_sampler_states(iteration=0)
    xunit = sampler_states[0].positions[0].unit

    for state_index in range(len(states)):
        state_positions = extract_trajectory(topology, state_index=state_index,
            output_data=output_data_path, checkpoint_data=checkpoint_data,
            frame_begin=frame_begin, frame_stride=frame_stride)

        file_name = f"{output_dir}/state_{state_index+1}.pdb"

        # Trajectories are written in nanometers:
        state_traj = Trajectory(
            state_positions,
            Topology.from_openmm(topology),
        )

        if center:
            ref_traj = state_traj[0]
            state_traj.superpose(ref_traj)
            # This rewrites to state_traj

        Trajectory.save_pdb(state_traj,file_name)

    return file_list

def extract_trajectory(
    topology, output_data="output/output.nc", checkpoint_data="output_checkpoint.nc",
    state_index=None, replica_index=None,
    frame_begin=0, frame_stride=1, frame_end=-1):
    """
    Internal function for extract trajectory (replica or state) from .nc file,
    Based on YANK extract_trajectory code.
    """

    reporter = MultiStateReporter(output_data, open_mode='r', checkpoint_storage=checkpoint_data)

    # Get dimensions
    trajectory_storage = reporter._storage_checkpoint
    n_iterations = reporter.read_last_iteration()
    n_frames = trajectory_storage.variables['positions'].shape[0]
    n_atoms = trajectory_storage.variables['positions'].shape[2]

    # Determine frames to extract.
    # Convert negative indices to last indices.
    if frame_begin < 0:
        frame_begin = n_frames + frame_begin
    if frame_end < 0:
        frame_end = n_frames + frame_end + 1
    frame_indices = range(frame_begin, frame_end, frame_stride)
    if len(frame_indices) == 0:
        raise ValueError('No frames selected')

    # Determine the number of frames that the trajectory will have.
    if state_index is None:
        n_trajectory_frames = len(frame_indices)
    else:
        # With SAMS, an iteration can have 0 or more replicas in a given state.
        # Deconvolute state indices.
        state_indices = [None for _ in frame_indices]
        for i, iteration in enumerate(frame_indices):
            replica_indices = reporter._storage_analysis.variables['states'][iteration, :]
            state_indices[i] = np.where(replica_indices == state_index)[0]
        n_trajectory_frames = sum(len(x) for x in state_indices)

    # Initialize positions and box vectors arrays.
    # MDTraj Cython code expects float32 positions.
    positions = np.zeros((n_trajectory_frames, n_atoms, 3), dtype=np.float32)

    # Extract state positions and box vectors.
    if state_index is not None:
        # Extract state positions
        frame_idx = 0
        for i, iteration in enumerate(frame_indices):
            for replica_index in state_indices[i]:
                positions[frame_idx, :, :] = trajectory_storage.variables['positions'][iteration, replica_index, :, :].astype(np.float32)
                frame_idx += 1

    else:  # Extract replica positions
        for i, iteration in enumerate(frame_indices):
            positions[i, :, :] = trajectory_storage.variables['positions'][iteration, replica_index, :, :].astype(np.float32)

    return positions

def process_replica_exchange_data(
    output_data="output/output.nc", output_directory="output", series_per_page=4,
    write_data_file=True, plot_production_only=False, print_timing=False,
    equil_nskip=1, frame_begin=0, frame_end=-1,
):
    """
    Read replica exchange simulation data, detect equilibrium and decorrelation time, and plot replica exchange results.

    :param output_data: path to output .nc file from replica exchange simulation, (default='output/output.nc')
    :type output_data: str

    :param output_directory: path to which output files will be written (default='output')
    :type output_directory: stry

    :param series_per_page: number of replica data series to plot per pdf page (default=4)
    :type series_per_page: int

    :param write_data_file: Option to write a text data file containing the state_energies array (default=True)
    :type write_data_file: Boolean

    :param plot_production_only: Option to plot only the production region, as determined from pymbar detectEquilibration (default=False)
    :type plot_production_only: Boolean

    :param equil_nskip: skip this number of frames to sparsify the energy timeseries for pymbar detectEquilibration (default=1) - this is used only when frame_begin=0 and the trajectory has less than 40000 frames.
    :type equil_nskip: Boolean

    :param frame_begin: analyze starting from this frame, discarding all prior as equilibration period (default=0)
    :type frame_begin: int

    :param frame_end: analyze up to this frame only, discarding the rest (default=-1).
    :type frame_end: int

    :returns:
        - replica_energies ( `Quantity() <http://docs.openmm.org/development/api-python/generated/simtk.unit.quantity.Quantity.html>`_ ( np.float( [number_replicas,number_simulation_steps] ), simtk.unit ) ) - The potential energies for all replicas at all (printed) time steps
        - replica_state_indices ( np.int64( [number_replicas,number_simulation_steps] ), simtk.unit ) - The thermodynamic state assignments for all replicas at all (printed) time steps
        - production_start ( int - The frame at which the production region begins for all replicas, as determined from pymbar detectEquilibration
        - sample_spacing ( int - The number of frames between uncorrelated state energies, estimated using heuristic algorithm )
        - n_transit ( np.float( [number_replicas] ) ) - Number of half-transitions between state 0 and n for each replica
        - mixing_stats ( tuple ( np.float( [number_replicas x number_replicas] ) , np.float( [ number_replicas ] ) , float( statistical inefficiency ) ) ) - transition matrix, corresponding eigenvalues, and statistical inefficiency
    """

    t1 = time.perf_counter()

    # Read the simulation coordinates for individual temperature replicas
    reporter = MultiStateReporter(output_data, open_mode="r")

    t2 = time.perf_counter()
    if print_timing:
        print(f"open data time: {t2-t1}")

    # figure out what the time between output is.
    # We assume all use the same time step (which i think is required)

    mcmove = reporter.read_mcmc_moves()[0]
    time_interval = mcmove.n_steps*mcmove.timestep

    t3 = time.perf_counter()
    if print_timing:
        print(f"read_mcmc_moves time: {t3-t2}")

    # figure out what the temperature list is
    states = reporter.read_thermodynamic_states()[0]

    t4 = time.perf_counter()
    if print_timing:
        print(f"read_thermodynamics_states time: {t4-t3}")

    temperature_list = []
    for s in states:
        temperature_list.append(s.temperature)

    analyzer = ReplicaExchangeAnalyzer(reporter)

    t5 = time.perf_counter()

    (
        replica_energies,
        unsampled_state_energies,
        neighborhoods,
        replica_state_indices,
    ) = analyzer.read_energies()

    # Truncate output of read_energies() to last frame of interest
    if frame_end > 0:
        # Use frames from frame_begin to frame_end
        replica_energies = replica_energies[:,:,:frame_end]
        unsampled_state_energies = unsampled_state_energies[:,:,:frame_end]
        neighborhoods = neighborhoods[:,:,:frame_end]
        replica_state_indices = replica_state_indices[:,:frame_end]

    t6 = time.perf_counter()
    if print_timing:
        print(f"read_energies time: {t6-t5}")

    n_particles = np.shape(reporter.read_sampler_states(iteration=0)[0].positions)[0]
    temps = np.array([temp._value for temp in temperature_list])
    beta_k = 1 / (kB * temps)
    n_replicas = len(temperature_list)
    for k in range(n_replicas):
        replica_energies[:, k, :] *= beta_k[k] ** (-1)

    t7 = time.perf_counter()
    if print_timing:
        print(f"reduce replica energies time: {t7-t6}")

    total_steps = len(replica_energies[0][0])
    state_energies = np.zeros([n_replicas, total_steps])

    t8 = time.perf_counter()
    # there must be some better way to do this as list comprehension.
    for step in range(total_steps):
        for state in range(n_replicas):
            state_energies[state, step] = replica_energies[
                np.where(replica_state_indices[:, step] == state)[0], 0, step
            ]

    t9 = time.perf_counter()
    if print_timing:
        print(f"assign state energies time: {t9-t8}")

    # can run physical-valication on these state_energies

    # Use pymbar timeseries module to detect production period

    t10 = time.perf_counter()

    # Start of equilibrated data:
    t0 = np.zeros((n_replicas))
    # Statistical inefficiency:
    g = np.zeros((n_replicas))

    subsample_indices = {}

    # If sufficiently large, discard the first 20000 frames as equilibration period and use
    # subsampleCorrelatedData to get the energy decorrelation time.
    if total_steps >= 40000 or frame_begin > 0:
        if frame_begin > 0:
            # If specified, use frame_begin as the start of the production region
            production_start=frame_begin
        else:
            # Otherwise, use frame 20000
            production_start=20000

        for state in range(n_replicas):
            subsample_indices[state] = timeseries.subsampleCorrelatedData(
                state_energies[state][production_start:],
                conservative=True,
            )
            g[state] = subsample_indices[state][1]-subsample_indices[state][0]

    else:
        # For small trajectories, use detectEquilibration
        for state in range(n_replicas):
            t0[state], g[state], Neff_max = timeseries.detectEquilibration(state_energies[state], nskip=equil_nskip)

            # Choose the latest equil timestep to apply to all states
            production_start = int(np.max(t0))

    # Assume a normal distribution (very rough approximation), and use mean plus
    # the number of standard deviations which leads to (n_replica-1)/n_replica coverage
    # For 12 replicas this should be the mean + 1.7317 standard deviations

    # x standard deviations is the solution to (n_replica-1)/n_replica = erf(x/sqrt(2))
    # This is equivalent to a target of 23/24 CDF value

    print(f"g: {g.astype(int)}")

    def erf_fun(x):
        return np.power((erf(x/np.sqrt(2))-(n_replicas-1)/n_replicas),2)

    # x must be larger than zero
    opt_g_results = minimize_scalar(
        erf_fun,
        bounds=(0,10)
        )

    if not opt_g_results.success:
        print("Error solving for correlation time, exiting...")
        print(f"erf opt results: {opt_g_results}")
        exit()

    sample_spacing = int(np.ceil(np.mean(g)+opt_g_results.x*np.std(g)))

    t11 = time.perf_counter()
    if print_timing:
        print(f"detect equil and subsampling time: {t11-t10}")

    print("state    mean energies  variance")
    for state in range(n_replicas):
        state_mean = np.mean(state_energies[state,production_start::sample_spacing])
        state_std = np.std(state_energies[state,production_start::sample_spacing])
        print(
            f"  {state:4d}    {state_mean:10.6f} {state_std:10.6f}"
        )

    t12 = time.perf_counter()

    if write_data_file == True:
        f = open(os.path.join(output_directory, "replica_energies.dat"), "w")
        for step in range(total_steps):
            f.write(f"{step:10d}")
            for replica_index in range(n_replicas):
                f.write(f"{replica_energies[replica_index,replica_index,step]:12.6f}")
            f.write("\n")
        f.close()

    t13 = time.perf_counter()
    if print_timing:
        print(f"Optionally write .dat file: {t13-t12}")

    t14 = time.perf_counter()

    if plot_production_only==True:
        plot_replica_exchange_energies(
            state_energies[:,production_start:],
            temperature_list,
            series_per_page,
            time_interval=time_interval,
            time_shift=production_start*time_interval,
            file_name=f"{output_directory}/rep_ex_ener.pdf",
        )

        plot_replica_exchange_energy_histograms(
            state_energies[:,production_start:],
            temperature_list,
            file_name=f"{output_directory}/rep_ex_ener_hist.pdf",
        )

        plot_replica_exchange_summary(
            replica_state_indices[:,production_start:],
            temperature_list,
            series_per_page,
            time_interval=time_interval,
            time_shift=production_start*time_interval,
            file_name=f"{output_directory}/rep_ex_states.pdf",
        )

        plot_replica_state_matrix(
            replica_state_indices[:,production_start:],
            file_name=f"{output_directory}/state_probability_matrix.pdf",
        )

    else:
        plot_replica_exchange_energies(
            state_energies,
            temperature_list,
            series_per_page,
            time_interval=time_interval,
            file_name=f"{output_directory}/rep_ex_ener.pdf",
        )

        plot_replica_exchange_energy_histograms(
            state_energies,
            temperature_list,
            file_name=f"{output_directory}/rep_ex_ener_hist.pdf",
        )

        plot_replica_exchange_summary(
            replica_state_indices,
            temperature_list,
            series_per_page,
            time_interval=time_interval,
            file_name=f"{output_directory}/rep_ex_states.pdf",
        )

        plot_replica_state_matrix(
            replica_state_indices,
            file_name=f"{output_directory}/state_probability_matrix.pdf",
        )

    t15 = time.perf_counter()

    if print_timing:
        print(f"plotting time: {t15-t14}")

    # Analyze replica exchange state transitions
    # For each replica, how many times does the thermodynamic state go between state 0 and state n
    # For consistency with the other mixing statistics, use only the production region here

    replica_state_indices_prod = replica_state_indices[:,production_start:]

    # Number of one-way transitions from states 0 to n or states n to 0
    n_transit = np.zeros((n_replicas,1))

    # Replica_state_indices is [n_replicas x n_iterations]
    for rep in range(n_replicas):
        last_bound = None
        for i in range(replica_state_indices_prod.shape[1]):
            if replica_state_indices_prod[rep,i] == 0 or replica_state_indices_prod[rep,i] == (n_replicas-1):
                if last_bound is None:
                    # This is the first time state 0 or n is visited
                    pass
                else:
                    if last_bound != replica_state_indices_prod[rep,i]:
                        # This is a completed transition from 0 to n or n to 0
                        n_transit[rep] += 1
                last_bound = replica_state_indices_prod[rep,i]

    t16 = time.perf_counter()

    if print_timing:
        print(f"replica transition analysis: {t16-t15}")

    # Compute transition matrix from the analyzer
    mixing_stats = analyzer.generate_mixing_statistics(number_equilibrated=production_start)

    t17 = time.perf_counter()

    if print_timing:
        print(f"compute transition matrix: {t17-t16}")
        print(f"total time elapsed: {t17-t1}")

    return (replica_energies, replica_state_indices, production_start, sample_spacing, n_transit, mixing_stats)

def run_replica_exchange(
    topology,
    system,
    positions,
    total_simulation_time=1.0 * unit.picosecond,
    simulation_time_step=None,
    temperature_list=None,
    friction=1.0 / unit.picosecond,
    minimize=True,
    exchange_frequency=1000,
    output_data="output/output.nc",
    ):

    """
    Run a OpenMMTools replica exchange simulation using an OpenMM coarse grained model.

    :param topology: OpenMM Topology
    :type topology: `Topology() <https://simtk.org/api_docs/openmm/api4_1/python/classsimtk_1_1openmm_1_1app_1_1topology_1_1Topology.html>`_

    :param system: OpenMM System()
    :type system: `System() <https://simtk.org/api_docs/openmm/api4_1/python/classsimtk_1_1openmm_1_1openmm_1_1System.html>`_

    :param positions: Positions array for the model we would like to test
    :type positions: `Quantity() <http://docs.openmm.org/development/api-python/generated/simtk.unit.quantity.Quantity.html>`_ ( np.array( [cgmodel.num_beads,3] ), simtk.unit )

    :param total_simulation_time: Total run time for individual simulations
    :type total_simulation_time: `SIMTK <https://simtk.org/>`_ `Unit() <http://docs.openmm.org/7.1.0/api-python/generated/simtk.unit.unit.Unit.html>`_

    :param simulation_time_step: Simulation integration time step
    :type simulation_time_step: `SIMTK <https://simtk.org/>`_ `Unit() <http://docs.openmm.org/7.1.0/api-python/generated/simtk.unit.unit.Unit.html>`_

    :param temperature_list: List of temperatures for which to perform replica exchange simulations, default = None
    :type temperature: List( float * simtk.unit.temperature )

    :param friction: Langevin thermostat friction coefficient, default = 1 / ps
    :type friction: `SIMTK <https://simtk.org/>`_ `Unit() <http://docs.openmm.org/7.1.0/api-python/generated/simtk.unit.unit.Unit.html>`_

    :param minimize: Whether minimization is done before running the simulation
    :type minimize: bool

    :param output_data: Name of NETCDF file where we will write simulation data
    :type output_data: string

    :param exchange_frequency: Number of time steps between replica exchange attempts, Default = None
    :type exchange_frequency: int

    :param output_data: file to put the output .nc
    :type output_data: netCDF4 file as generated by OpenMM

    :returns:
        - replica_energies ( `Quantity() <http://docs.openmm.org/development/api-python/generated/simtk.unit.quantity.Quantity.html>`_ ( np.float( [number_replicas,number_simulation_steps] ), simtk.unit ) ) - The potential energies for all replicas at all (printed) time steps
        - replica_positions ( `Quantity() <http://docs.openmm.org/development/api-python/generated/simtk.unit.quantity.Quantity.html>`_ ( np.float( [number_replicas,number_simulation_steps,cgmodel.num_beads,3] ), simtk.unit ) ) - The positions for all replicas at all (printed) time steps

        - replica_state_indices ( np.int64( [number_replicas,number_simulation_steps] ), simtk.unit ) - The thermodynamic state assignments for all replicas at all (printed) time steps

    :Example:

    >>> from foldamers.cg_model.cgmodel import CGModel
    >>> from cg_openmm.simulation.rep_exch import *
    >>> cgmodel = CGModel()
    >>> replica_energies,replica_positions,replica_state_indices = run_replica_exchange(cgmodel.topology,cgmodel.system,cgmodel.positions)

    """

    simulation_steps = int(np.floor(total_simulation_time / simulation_time_step))

    exchange_attempts = int(np.floor(simulation_steps / exchange_frequency))

    if temperature_list is None:
        temperature_list = [((300.0 + i) * unit.kelvin) for i in range(-50, 50, 10)]

    num_replicas = len(temperature_list)
    sampler_states = list()
    thermodynamic_states = list()

    # Define thermodynamic states.
    # box_vectors = system.getDefaultPeriodicBoxVectors()
    for temperature in temperature_list:
        thermodynamic_state = openmmtools.states.ThermodynamicState(
            system=system, temperature=temperature
        )
        thermodynamic_states.append(thermodynamic_state)
        sampler_states.append(
            openmmtools.states.SamplerState(positions)
        )  # no box vectors, non-periodic system.

    # Create and configure simulation object.
    move = openmmtools.mcmc.LangevinDynamicsMove(
        timestep=simulation_time_step,
        collision_rate=friction,
        n_steps=exchange_frequency,
        reassign_velocities=False,
    )

    simulation = ReplicaExchangeSampler(
        mcmc_moves=move,
        number_of_iterations=exchange_attempts,
        replica_mixing_scheme='swap-neighbors',
    )

    if os.path.exists(output_data):
        os.remove(output_data)

    reporter = MultiStateReporter(output_data, checkpoint_interval=1)
    simulation.create(thermodynamic_states, sampler_states, reporter)

    if minimize:
        simulation.minimize()

    print("Running OpenMM replica exchange simulation...")
    print(f"Time step: {simulation_time_step}")
    print(f"Iterations: {exchange_attempts}")
    try:
        simulation.run()
    except BaseException:
        print("Replica exchange simulation failed, try verifying your model/simulation settings.")
        exit()

    return

def restart_replica_exchange(
    total_simulation_time=1*unit.nanosecond,
    simulation_time_step=5*unit.picosecond,
    exchange_frequency=200,
    output_data="output/output.nc",
    ):

    """
    Restart an OpenMMTools replica exchange simulation using an OpenMM coarse grained model and
    output .nc files from the previous segment of the simulation.

    :param total_simulation_time: Total run time to add to the original simulation (default=1*unit.nanosecond)
    :type total_simulation_time: `SIMTK <https://simtk.org/>`_ `Unit() <http://docs.openmm.org/7.1.0/api-python/generated/simtk.unit.unit.Unit.html>`_

    :param simulation_time_step: Simulation integration time step (default=5*unit.picosecond)
    :type simulation_time_step: `SIMTK <https://simtk.org/>`_ `Unit() <http://docs.openmm.org/7.1.0/api-python/generated/simtk.unit.unit.Unit.html>`_

    :param exchange_frequency: Number of time steps between replica exchange attempts (default=200)
    :type exchange_frequency: int

    :param output_data: Path to the NETCDF file for previous segment of simulation - this will be appended to (default="output/output.nc")
    :type output_data: str
    """

    simulation_steps = int(np.floor(total_simulation_time / simulation_time_step))
    exchange_attempts = int(np.floor(simulation_steps / exchange_frequency))

    # Load in the reporter from the original simulation:
    reporter = MultiStateReporter(output_data, open_mode="r+")
    simulation = ReplicaExchangeSampler.from_storage(reporter)

    print("Running OpenMM replica exchange simulation...")
    print(f"Time step: {simulation_time_step}")
    print(f"Iterations: {exchange_attempts}")

    simulation.extend(n_iterations=exchange_attempts)

    return

def get_minimum_energy_ensemble(
    topology, replica_energies, replica_positions, ensemble_size=5, file_name=None
):

    """
    Get an ensemble of low (potential) energy poses, and write the lowest energy structure to a PDB file if a file_name is provided.

    :param topology: OpenMM Topology()
    :type topology: `Topology() <https://simtk.org/api_docs/openmm/api4_1/python/classsimtk_1_1openmm_1_1app_1_1topology_1_1Topology.html>`_

    :param replica_energies: List of dimension num_replicas X simulation_steps, which gives the energies for all replicas at all simulation steps
    :type replica_energies: List( List( float * simtk.unit.energy for simulation_steps ) for num_replicas )

    :param replica_positions: List of positions for all output frames for all replicas
    :type replica_positions: np.array( ( float * simtk.unit.positions for num_beads ) for simulation_steps )

    :param file_name: Output destination for PDB coordinates of minimum energy pose, Default = None

    :returns:
    - ensemble ( List() ) - A list of poses that are in the minimum energy ensemble.

    :Example:

    >>> from foldamers.cg_model.cgmodel import CGModel
    >>> from cg_openmm.simulation.rep_exch import *
    >>> cgmodel = CGModel()
    >>> replica_energies,replica_positions,replica_state_indices = run_replica_exchange(cgmodel.topology,cgmodel.system,cgmodel.positions)
    >>> ensemble_size = 5
    >>> file_name = "minimum.pdb"
    >>> minimum_energy_ensemble = get_minimum_energy_ensemble(cgmodel.topology,replica_energies,replica_positions,ensemble_size=ensemble_size,file_name=file_name)

    """
    # Get the minimum energy structure sampled during the simulation
    ensemble = []
    ensemble_energies = []
    for replica in range(len(replica_energies)):
        energies = np.array([energy for energy in replica_energies[replica][replica]])
        for energy in range(len(energies)):
            if len(ensemble) < ensemble_size:
                ensemble.append(replica_positions[replica][energy])
                ensemble_energies.append(energies[energy])
            else:
                for comparison in range(len(ensemble_energies)):
                    if energies[energy] < ensemble_energies[comparison]:
                        ensemble_energies[comparison] = energies[energy]
                        ensemble[comparison] = replica_positions[replica][energy]

    if file_name is None:
        index = 1
        for pose in ensemble:
            file = open(str("re_min_" + str(index) + ".pdb"), "w")
            PDBFile.writeFile(topology, pose, file=file)
    else:
        file = open(file_name, "w")
        for pose in ensemble:
            PDBFile.writeFile(topology, pose, file=file)

    return ensemble

def plot_replica_exchange_energies(
    state_energies,
    temperature_list,
    series_per_page,
    time_interval=1.0 * unit.picosecond,
    time_shift=0.0 * unit.picosecond,
    file_name="rep_ex_ener.pdf",
    legend=True,
):
    """
    Plot the potential energies for a batch of replica exchange trajectories

    :param state_energies: List of dimension num_replicas X simulation_steps, which gives the energies for all replicas at all simulation steps
    :type state_energies: List( List( float * simtk.unit.energy for simulation_steps ) for num_replicas )

    :param temperature_list: List of temperatures for which to perform replica exchange simulations, default = [(300.0 * unit.kelvin).__add__(i * unit.kelvin) for i in range(-20,100,10)]
    :type temperature: List( float * simtk.unit.temperature )

    :param time_interval: interval between energy exchanges.
    :type time_interval: `SIMTK <https://simtk.org/>`_ `Unit() <http://docs.openmm.org/7.1.0/api-python/generated/simtk.unit.unit.Unit.html>`_

    :param time_shift: amount of time before production period to shift the time axis(default = 0)
    :type time_shift: `SIMTK <https://simtk.org/>`_ `Unit() <http://docs.openmm.org/7.1.0/api-python/generated/simtk.unit.unit.Unit.html>`_

    :param file_name: The pathname of the output file for plotting results, default = "replica_exchange_energies.png"
    :type file_name: str

    :param legend: Controls whether a legend is added to the plot
    :type legend: Logical

    """

    simulation_times = np.array(
        [
            step * time_interval.value_in_unit(unit.picosecond)
            for step in range(len(state_energies[0]))
        ]
    )

    simulation_times += time_shift.value_in_unit(unit.picosecond)

    # To improve pdf render speed, sparsify data to display less than 2000 data points
    n_xdata = len(simulation_times)

    if n_xdata <= 1000:
        plot_stride = 1
    else:
        plot_stride = int(np.floor(n_xdata/1000))

    # If more than series_per_page replicas, split into separate pages for better visibility
    nmax = series_per_page
    npage = int(np.ceil(len(temperature_list)/nmax))

    with PdfPages(file_name) as pdf:
        page_num=1
        plotted_per_page=0
        pyplot.figure()
        for state in range(len(temperature_list)):
            if plotted_per_page <= (nmax):
                pyplot.plot(
                    simulation_times[::plot_stride],
                    state_energies[state,::plot_stride],
                    alpha=0.5,
                    linewidth=1,
                )
                plotted_per_page += 1

            if (plotted_per_page >= nmax) or (state==(len(temperature_list)-1)):
                # Save and close previous page
                pyplot.xlabel("Simulation Time ( Picoseconds )")
                pyplot.ylabel("Potential Energy ( kJ / mol )")
                pyplot.title("Replica Exchange Simulation")

                if legend:
                    pyplot.legend(
                        [round(temperature.value_in_unit(unit.kelvin), 1) for temperature in temperature_list[(0+(page_num-1)*nmax):(page_num*nmax)]],
                        loc="center left",
                        bbox_to_anchor=(1, 0.5),
                        title="T (K)",
                    )

                pdf.savefig(bbox_inches="tight") # Save current fig to pdf page
                pyplot.close()
                plotted_per_page = 0
                page_num += 1

    return

def plot_replica_exchange_energy_histograms(
    state_energies,
    temperature_list,
    file_name="rep_ex_ener_hist.pdf",
    legend=True,
):
    """
    Plot the potential energies for a batch of replica exchange trajectories

    :param state_energies: List of dimension num_replicas X simulation_steps, which gives the energies for all replicas at all simulation steps
    :type state_energies: List( List( float * simtk.unit.energy for simulation_steps ) for num_replicas )

    :param temperature_list: List of temperatures for which to perform replica exchange simulations, default = [(300.0 * unit.kelvin).__add__(i * unit.kelvin) for i in range(-20,100,10)]
    :type temperature: List( float * simtk.unit.temperature )

    :param file_name: The pathname of the output file for plotting results, default = "replica_exchange_energies.png"
    :type file_name: str

    :param legend: Controls whether a legend is added to the plot
    :type legend: Logical

    """

    figure = pyplot.figure(figsize=(8.5,11))

    for state in range(len(temperature_list)):
        n_out, bin_edges_out = np.histogram(
            state_energies[state,:],bins=20,density=True,
        )

        bin_centers = np.zeros((len(bin_edges_out)-1,1))
        for i in range(len(bin_edges_out)-1):
            bin_centers[i] = (bin_edges_out[i]+bin_edges_out[i+1])/2

        pyplot.plot(bin_centers,n_out,'o-',alpha=0.5,linewidth=1,markersize=6)

    pyplot.xlabel("Potential Energy ( kJ / mol )")
    pyplot.ylabel("Probability")
    pyplot.title("Replica Exchange Energy Histogram")

    if legend:
        pyplot.legend(
            [round(temperature._value, 1) for temperature in temperature_list],
            loc="center left",
            bbox_to_anchor=(1, 0.5),
            title="T (K)",
        )

    pyplot.savefig(file_name, bbox_inches="tight")
    pyplot.close()

    return

def plot_replica_state_matrix(
    replica_state_indices,
    file_name='state_probability_matrix.pdf'
    ):

    # Plot a matrix of replica vs. state, coloring each box in the grid by normalized frequency
    # For each replica, histogram the state indices data
    # Then normalize the data and create [n_replica x n_state] patch graph

    n_replicas = replica_state_indices.shape[0]

    hist_all = np.zeros((n_replicas, n_replicas))

    state_bin_edges = np.linspace(-0.5,n_replicas-0.5,n_replicas+1)
    state_bin_centers = 0.5+state_bin_edges[0:n_replicas]

    for rep in range(n_replicas):
        hist_all[rep,:], bin_edges = np.histogram(
            replica_state_indices[rep,:],bins=state_bin_edges,density=True,
        )

    # No need for global normalization, since each replica's state probabilities must sum to 1

    hist_norm = np.zeros_like(hist_all)
    for rep in range(n_replicas):
        for state in range(n_replicas):
            hist_norm[rep,state] = hist_all[rep,state]/np.max(hist_all[rep,:])

    mean_score = np.mean(hist_norm)
    min_score = np.amin(hist_norm)

    ax = pyplot.subplot(111)

    cmap=pyplot.get_cmap('nipy_spectral')
    norm=Normalize(vmin=0,vmax=1)

    ax.imshow(hist_norm,cmap=cmap,norm=norm)
    ax.set_aspect('equal', 'box')

    # Append colorbar axis to right side
    divider = make_axes_locatable(ax)
    cax = divider.append_axes("right",size="5%",pad=0.20)

    pyplot.colorbar(
        cm.ScalarMappable(cmap=cmap,norm=norm),
        cax=cax,
        label='normalized frequency',
        )

    ax.set_xlabel("State")
    ax.set_ylabel("Replica")
    pyplot.suptitle(f"Replica exchange state probabilities\n(Mean: {mean_score:.4f} Min: {min_score:.4f})")

    pyplot.savefig(file_name)
    pyplot.close()

    return hist_all

def plot_replica_exchange_summary(
    replica_states,
    temperature_list,
    series_per_page,
    time_interval=1.0 * unit.picosecond,
    time_shift=0.0 * unit.picosecond,
    file_name="rep_ex_states.pdf",
    legend=True,
):
    """
    Plot the thermodynamic state assignments for individual temperature replicas as a function of the simulation time, in order to obtain a visual summary of the replica exchanges from a OpenMM simulation.

    :param replica_states: List of dimension num_replicas X simulation_steps, which gives the thermodynamic state indices for all replicas at all simulation steps
    :type replica_states: List( List( float * simtk.unit.energy for simulation_steps ) for num_replicas )

    :param temperature_list: List of temperatures for which to perform replica exchange simulations, default = [(300.0 * unit.kelvin).__add__(i * unit.kelvin) for i in range(-20,100,10)]
    :type temperature: List( float * simtk.unit.temperature )

    :param time_interval: interval between energy exchanges.
    :type time_interval: `SIMTK <https://simtk.org/>`_ `Unit() <http://docs.openmm.org/7.1.0/api-python/generated/simtk.unit.unit.Unit.html>`_

    :param time_shift: amount of time before production period to shift the time axis(default = 0)
    :type time_shift: `SIMTK <https://simtk.org/>`_ `Unit() <http://docs.openmm.org/7.1.0/api-python/generated/simtk.unit.unit.Unit.html>`_

    :param file_name: The pathname of the output file for plotting results, default = "replica_exchange_state_transitions.png"
    :type file_name: str

    :param legend: Controls whether a legend is added to the plot
    :type legend: Logical

    """

    simulation_times = np.array(
        [
            step * time_interval.value_in_unit(unit.picosecond)
            for step in range(len(replica_states[0]))
        ]
    )

    simulation_times += time_shift.value_in_unit(unit.picosecond)

    # To improve pdf render speed, sparsify data to display less than 2000 data points
    n_xdata = len(simulation_times)

    if n_xdata <= 1000:
        plot_stride = 1
    else:
        plot_stride = int(np.floor(n_xdata/1000))

    # If more than series_per_page replicas, split into separate pages for better visibility
    nmax = series_per_page
    npage = int(np.ceil(len(temperature_list)/nmax))

    with PdfPages(file_name) as pdf:
        page_num=1
        plotted_per_page=0
        pyplot.figure()
        for replica in range(len(replica_states)):
            state_indices = np.array([int(round(state)) for state in replica_states[replica]])

            if plotted_per_page <= (nmax):

                pyplot.plot(
                    simulation_times[::plot_stride],
                    state_indices[::plot_stride],
                    alpha=0.5,
                    linewidth=1
                )
                plotted_per_page += 1

            if (plotted_per_page >= nmax) or (replica==(len(replica_states)-1)):
                # Save and close previous page
                pyplot.xlabel("Simulation Time ( Picoseconds )")
                pyplot.ylabel("Thermodynamic State Index")
                pyplot.title("State Exchange Summary")

                if legend:
                    pyplot.legend(
                        [i for i in range((page_num-1)*nmax,page_num*nmax)],
                        loc="center left",
                        bbox_to_anchor=(1, 0.5),
                        title="Replica Index",
                    )

                pdf.savefig(bbox_inches="tight") # Save current fig to pdf page
                pyplot.close()
                plotted_per_page = 0
                page_num += 1

    return


#!/usr/bin/env python3

"""
corrections.py: Script to apply corrections to the images.
"""

import os
from argparse import ArgumentParser
from datetime import date, datetime
from typing import Optional, Sequence

import numpy as np
from astropy.io import fits

from dresscode.utils import load_config

def main(argv: Optional[Sequence[str]] = None) -> int:
    parser = ArgumentParser()
    parser.add_argument(
        "-c", "--config", help="path to config.txt", default="config.txt"
    )
    args = parser.parse_args(argv)

    config = load_config(args.config)

    galaxy = config["galaxy"]
    path = config["path"] + galaxy + "/working_dir/"
    years = config["years"]

    # Loop over the different years.
    for year in years:

        print("Year: " + year)
        yearpath = path + year + "/"

        # PART 1: Apply a coincidence loss correction.

        print("Applying coincidence loss corrections...")
        if os.path.isfile(yearpath + "sum_um2_nm.img"):
            coicorr(yearpath + "sum_um2_nm.img")
        if os.path.isfile(yearpath + "sum_uw2_nm.img"):
            coicorr(yearpath + "sum_uw2_nm.img")
        if os.path.isfile(yearpath + "sum_uw1_nm.img"):
            coicorr(yearpath + "sum_uw1_nm.img")

        # PART 2: Apply a large scale sensitivity correction.

        print("Applying large scale sensitivity corrections...")
        if os.path.isfile(yearpath + "sum_um2_nm_coi.img"):
            lsscorr(yearpath + "sum_um2_nm_coi.img")
        if os.path.isfile(yearpath + "sum_uw2_nm_coi.img"):
            lsscorr(yearpath + "sum_uw2_nm_coi.img")
        if os.path.isfile(yearpath + "sum_uw1_nm_coi.img"):
            lsscorr(yearpath + "sum_uw1_nm_coi.img")

        # PART 3: Apply a zero point correction.

        print("Applying zero point corrections...")
        if os.path.isfile(yearpath + "sum_um2_nm_coilss.img"):
            zeropoint(yearpath + "sum_um2_nm_coilss.img", -2.330e-3, -1.361e-3)
        if os.path.isfile(yearpath + "sum_uw2_nm_coilss.img"):
            zeropoint(yearpath + "sum_uw2_nm_coilss.img", 1.108e-3, -1.960e-3)
        if os.path.isfile(yearpath + "sum_uw1_nm_coilss.img"):
            zeropoint(yearpath + "sum_uw1_nm_coilss.img", 2.041e-3, -1.748e-3)

    return 0

# Functions for PART 1: Coincidence loss correction.
def coicorr(filename):
    # Open the image. Create arrays with zeros with the shape of the image.
    hdulist = fits.open(filename)
    data = hdulist[0].data
    header = hdulist[0].header
    total_flux = np.full_like(data, np.nan, dtype=np.float64)
    std = np.full_like(data, np.nan, dtype=np.float64)

    # Loop over all pixels and for each pixel: sum the flux densities (count rates) of
    # the 9x9 surrounding pixels: Craw (counts/s). Calculate the standard deviation in
    # the 9x9 pixels box.
    for x in range(5, data.shape[1] - 5):
        for y in range(5, data.shape[0] - 5):
            total_flux[y, x] = np.sum(data[y - 4 : y + 5, x - 4 : x + 5])
            std[y, x] = np.std(data[y - 4 : y + 5, x - 4 : x + 5])

    # Obtain the dead time correction factor and the frame time (in s) from the header
    # of the image.
    alpha = header["DEADC"]
    ft = header["FRAMTIME"]

    # Calculate the total number of counts in the 9x9 pixels box: x = Craw*ft (counts).
    # Calculate the minimum and maximum possible number of counts in the 9x9 pixels box.
    total_counts = ft * total_flux
    total_counts_min = ft * (total_flux - 81 * std)
    total_counts_max = ft * (total_flux + 81 * std)

    # Calculate the polynomial correction factor and the minimum and maximum possible
    # polynomial correction factor.
    f = polynomial(total_counts)
    f_min = polynomial(total_counts_min)
    f_max = polynomial(total_counts_max)

    # If alpha*total_counts_max is larger than 1, replace this value by 0.99. Otherwise,
    # the maximum possible theoretical coincidence-loss-corrected count rate will be NaN
    # in these pixels.
    if np.sum(alpha * total_counts_max >= 1.0) != 0:
        print(
            "Warning: The following pixels have very high fluxes. The uncertainty on "
            "the correction factor for these pixels is not to be trusted!",
            np.where(alpha * total_counts_max >= 1.0),
        )
    total_counts_max[alpha * total_counts_max >= 1.0] = 0.99 / alpha

    # Calculate the theoretical coincidence-loss-corrected count rate:
    # Ctheory = -ln(1 - alpha*Craw*ft) / (alpha*ft) (counts/s).
    # Calculate the minimum and maximum possible theoretical coincidence-loss-corrected
    # count rate.
    Ctheory = -np.log1p(-alpha * total_counts) / (alpha * ft)
    Ctheory_min = -np.log1p(-alpha * total_counts_min) / (alpha * ft)
    Ctheory_max = -np.log1p(-alpha * total_counts_max) / (alpha * ft)

    # Calculate the coincidence loss correction factor:
    # Ccorrfactor = Ctheory*f(x)/Craw.
    # Calculate the minimum and maximum possible coincidence loss correction factor.
    corrfactor = (Ctheory * f) / total_flux
    corrfactor_min = (Ctheory_min * f_min) / (total_flux - 81 * std)
    corrfactor_max = (Ctheory_max * f_max) / (total_flux + 81 * std)

    # Apply the coincidence loss correction to the data. Apply the minimum and maximum
    # coincidence loss correction to the data.
    new_data = corrfactor * data
    new_data_min = corrfactor_min * data
    new_data_max = corrfactor_max * data

    # Calculate the uncertainty and the relative uncertainty on the coincidence loss
    # correction. Put the relative uncertainty to 0 if the uncertainty is 0 (because in
    # those pixels the flux is also 0 and the relative uncertainty would be NaN).
    coicorr_unc = np.maximum(
        np.abs(new_data - new_data_min), np.abs(new_data_max - new_data)
    )
    coicorr_rel = coicorr_unc / new_data
    coicorr_rel[coicorr_unc == 0.0] = 0.0

    print(
        "The median coincidence loss correction factor for image "
        + os.path.basename(filename)
        + " is "
        + str(np.nanmedian(corrfactor))
        + " and the median relative uncertainty on the corrected data is "
        + str(np.nanmedian(coicorr_rel))
        + "."
    )

    # Adapt the header. Write the corrected data, the applied coincidence loss
    # correction and the relative uncertainty to a new image.
    header["PLANE0"] = "primary (counts/s)"
    header["PLANE1"] = "coincidence loss correction factor"
    header["PLANE2"] = "relative coincidence loss correction uncertainty (fraction)"
    datacube = [new_data, corrfactor, coicorr_rel]
    new_hdu = fits.PrimaryHDU(datacube, header)
    new_hdu.writeto(filename.replace(".img", "_coi.img"), overwrite=True)

    print(os.path.basename(filename) + " has been corrected for coincidence loss.")

# Function to calculate the empirical polynomial correction to account for the
# differences between the observed and theoretical coincidence loss correction:
# f(x) = 1 + a1x + a2x**2 + a3x**3 + a4x**4.
def polynomial(x):
    a1 = 0.0658568
    a2 = -0.0907142
    a3 = 0.0285951
    a4 = 0.0308063
    return 1 + (a1 * x) + (a2 * x ** 2) + (a3 * x ** 3) + (a4 * x ** 4)

# Function for PART 2: Large scale sensitivity correction.
def lsscorr(filename):
    # Open the image and the large scale sensitivity map.
    hdulist = fits.open(filename)
    data = hdulist[0].data[0]
    coicorr = hdulist[0].data[1]
    coicorr_rel = hdulist[0].data[2]
    header = hdulist[0].header

    lss_hdulist = fits.open(filename.replace("nm_coi", "lss"))
    lss_data = lss_hdulist[1].data

    # Apply the large scale sensitivity correction to the data.
    new_data = data / lss_data
    new_datacube = [new_data, coicorr, coicorr_rel]

    # Write the corrected data to a new image.
    new_hdu = fits.PrimaryHDU(new_datacube, header)
    new_hdu.writeto(filename.replace(".img", "lss.img"), overwrite=True)

    print(
        os.path.basename(filename)
        + " has been corrected for large scale sensitivity variations."
    )

# Function for PART 3: Zero point correction.
def zeropoint(filename, param1, param2):
    # Open the file.
    hdulist = fits.open(filename)
    data = hdulist[0].data[0]
    coicorr = hdulist[0].data[1]
    coicorr_rel = hdulist[0].data[2]
    header = hdulist[0].header
    # Calculate the average date of observation.
    start_date = datetime.strptime(header["DATE-OBS"].split("T")[0], "%Y-%m-%d").date()
    end_date = datetime.strptime(header["DATE-END"].split("T")[0], "%Y-%m-%d").date()
    obs_date = (end_date - start_date) / 2 + start_date
    # Calculate the number of years that have elapsed since the 1st of January 2005.
    first_date = date(2005, 1, 1)
    elapsed_time = obs_date - first_date
    years_passed = elapsed_time.days / 365.25

    # Calculate the zero point correction.
    zerocorr = 1 + param1 * years_passed + param2 * years_passed ** 2

    # Apply the correction to the data.
    new_data = data / zerocorr

    # Adapt the header. Write the corrected data to a new image.
    header["ZPCORR"] = zerocorr
    datacube = [new_data, coicorr, coicorr_rel]
    new_hdu = fits.PrimaryHDU(datacube, header)
    new_hdu.writeto(filename.replace(".img", "zp.img"), overwrite=True)

    print(
        os.path.basename(filename)
        + " has been corrected for sensitivity loss of the detector over time."
    )

if __name__ == "__main__":
    exit(main())