# !/usr/bin/env python
# -*- coding: utf-8 -*-
#
#    Project: Azimuthal integration
#             https://github.com/silx-kit/pyFAI
#
#    Copyright (C) 2013-2020 European Synchrotron Radiation Facility, Grenoble, France
#
#  Permission is hereby granted, free of charge, to any person obtaining a copy
#  of this software and associated documentation files (the "Software"), to deal
#  in the Software without restriction, including without limitation the rights
#  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
#  copies of the Software, and to permit persons to whom the Software is
#  furnished to do so, subject to the following conditions:
#  .
#  The above copyright notice and this permission notice shall be included in
#  all copies or substantial portions of the Software.
#  .
#  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
#  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
#  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
#  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
#  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
#  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
#  THE SOFTWARE.

__author__ = "Jérôme Kieffer"
__contact__ = "Jerome.Kieffer@ESRF.eu"
__license__ = "MIT"
__copyright__ = "European Synchrotron Radiation Facility, Grenoble, France"
__date__ = "06/07/2020"
__status__ = "development"

import logging
import threading
import os
import numpy
logger = logging.getLogger(__name__)
from math import ceil, floor
from . import detectors
from .opencl import ocl
if ocl:
    from .opencl import azim_lut as ocl_azim_lut
    from .opencl import azim_csr as ocl_azim_csr
else:
    ocl_azim_lut = ocl_azim_csr = None
from .third_party import six

try:
    from .ext import _distortion
    from .ext import sparse_utils
except ImportError:
    logger.debug("Backtrace", exc_info=True)
    logger.warning("Import _distortion cython implementation failed ... pure python version is terribly slow !!!")
    _distortion = None

try:
    from scipy.sparse import linalg, csr_matrix, identity
except IOError:
    logger.warning("Scipy is missing ... uncorrection will be handled the old way")
    linalg = None
else:
    import scipy
    v = tuple(int(i) for i in scipy.version.short_version.split(".") if i.isdigit())
    if v < (0, 11):
        logger.warning("Scipy is too old ... uncorrection will be handled the old way")
        linalg = None

def resize_image_2D_numpy(image, shape_in):
    "numpy implementation of resize_image_2D"
    new_img = numpy.zeros(shape_in, dtype=image.dtype)
    common_shape = [min(i, j) for i, j in zip(image.shape, shape_in)]
    new_img[:common_shape[0], :common_shape[1]] = image[:common_shape[0], :common_shape[1]]
    return new_img

if _distortion is None:
    resize_image_2D = resize_image_2D_numpy
else:
    from .ext._distortion import resize_image_2D

class Distortion(object):
    """
    This class applies a distortion correction on an image.

    New version compatible both with CSR and LUT...
    """
    def __init__(self, detector="detector", shape=None, resize=False, empty=0,
                 mask=None, method="CSR", device=None, workgroup=32):
        """
        :param detector: detector instance or detector name
        :param shape: shape of the output image
        :param resize: allow the output shape to be different from the input shape
        :param empty: value to be given for empty bins
        :param method: "lut" or "csr", the former is faster
        :param device: Name of the device: None for OpenMP, "cpu" or "gpu" or the id of the OpenCL device a 2-tuple of integer
        :param workgroup: workgroup size for CSR on OpenCL
        """
        self._shape_out = None
        if isinstance(detector, six.string_types):
            self.detector = detectors.detector_factory(detector)
        else:  # we assume it is a Detector instance
            self.detector = detector
        self.shape_in = self.detector.shape
        if mask is None:
            self.mask = numpy.ascontiguousarray(self.detector.mask, numpy.int8)
        elif mask is False:
            self.mask = numpy.zeros(self.detector.mask.shape, numpy.int8)
        else:
            self.mask = numpy.ascontiguousarray(mask, numpy.int8)
        self.resize = resize
        if shape is not None:
            self._shape_out = tuple([int(i) for i in shape])
        elif not self.resize:
            if self.detector.shape is not None:
                self._shape_out = self.detector.shape
            else:
                raise RuntimeError("You need to provide either the detector or its shape")

        self._sem = threading.Semaphore()
        self.bin_size = None
        self.max_size = None
        self.pos = None
        if not method:
            self.method = "lut"
        else:
            self.method = method.lower()
        if (self.detector.uniform_pixel and self.detector.IS_FLAT):
            csr = identity(numpy.prod(self.detector.shape),
                           dtype=numpy.float32,
                           format="csr")
            if self.detector.mask is not None:
                masked = numpy.where(self.detector.mask)
                csr[masked] = 0.0
            if self.method == "lut":
                self.lut = sparse_utils.CSR_to_LUT(csr.data, csr.indices, csr.indptr)
            else:
                self.lut = csr.data, csr.indices, csr.indptr
        else:
            #initialize the LUT later
            self.lut = None
        self.delta1 = self.delta2 = None  # max size of an pixel on a regular grid ...
        self.offset1 = self.offset2 = 0  # position of the first bin
        self.integrator = None
        self.empty = empty  # "dummy" value for empty bins
        self.device = device
        if not workgroup:
            self.workgroup = 1
        else:
            self.workgroup = int(workgroup)

    def __repr__(self):
        return os.linesep.join(["Distortion correction %s on device %s for detector shape %s:" % (self.method, self.device, self._shape_out),
                                self.detector.__repr__()])

    def reset(self, method=None, device=None, workgroup=None, prepare=True):
        """
        reset the distortion correction and re-calculate the look-up table

        :param method: can be "lut" or "csr", "lut" looks faster
        :param device: can be None, "cpu" or "gpu" or the id as a 2-tuple of integer
        :param worgroup: enforce the workgroup size for CSR.
        :param prepare: set to false to only reset and not re-initialize
        """
        with self._sem:
            self.max_size = None
            self.pos = None
            self.lut = None
            self.delta1 = self.delta2 = None
            self.offset1 = self.offset2 = 0
            self.integrator = None
            if method is not None:
                self.method = method.lower()
            if device is not None:
                self.device = device
            if workgroup is not None:
                self.workgroup = int(workgroup)
        if prepare:
            self.calc_init()

    @property
    def shape_out(self):
        """
        Calculate/cache the output shape

        :return: output shape
        """
        if self._shape_out is None:
            self.calc_pos()
        return self._shape_out

    def calc_pos(self, use_cython=True):
        """Calculate the pixel boundary position on the regular grid

        :return: pixel corner positions (in pixel units) on the regular grid
        :rtype: ndarray of shape (nrow, ncol, 4, 2)
        """
        if self.delta1 is None:
            with self._sem:
                if self.delta1 is None:
                    # TODO: implement equivalent in Cython
                    if _distortion and use_cython:
                        self.pos, self.delta1, self.delta2, shape_out, offset = _distortion.calc_pos(self.detector.get_pixel_corners(), self.detector.pixel1, self.detector.pixel2, self._shape_out)
                        if self._shape_out is None:
                            self.offset1, self.offset2 = offset
                            self._shape_out = shape_out
                    else:
                        pixel_size = numpy.array([self.detector.pixel1, self.detector.pixel2], dtype=numpy.float32)
                        # make it a 4D array
                        pixel_size.shape = 1, 1, 1, 2
                        pixel_size.strides = 0, 0, 0, pixel_size.strides[-1]
                        self.pos = self.detector.get_pixel_corners()[..., 1:] / pixel_size
                        if self._shape_out is None:
                            # if defined, it is probably because resize=False
                            corner_pos = self.pos.view()
                            corner_pos.shape = -1, 2
                            pos1_min, pos2_min = corner_pos.min(axis=0)
                            pos1_max, pos2_max = corner_pos.max(axis=0)
                            self._shape_out = (int(ceil(pos1_max - pos1_min)),
                                               int(ceil(pos2_max - pos2_min)))
                            self.offset1, self.offset2 = pos1_min, pos2_min
                        pixel_delta = self.pos.view()
                        pixel_delta.shape = -1, 4, 2
                        self.delta1, self.delta2 = ((numpy.ceil(pixel_delta.max(axis=1)) - numpy.floor(pixel_delta.min(axis=1))).max(axis=0)).astype(int)
        return self.pos

    def calc_size(self, use_cython=True):
        """Calculate the number of pixels falling into every single bin and

        :return: max of pixel falling into a single bin

        Considering the "half-CCD" spline from ID11 which describes a (1025,2048) detector,
        the physical location of pixels should go from:
        [-17.48634 : 1027.0543, -22.768829 : 2028.3689]
        We chose to discard pixels falling outside the [0:1025,0:2048] range with a lose of intensity
        """
        if self.pos is None:
            pos = self.calc_pos()
        else:
            pos = self.pos
        if self.max_size is None:
            with self._sem:
                if self.max_size is None:
                    if _distortion and use_cython:
                        self.bin_size = _distortion.calc_size(self.pos, self._shape_out, self.mask, (self.offset1, self.offset2))
                    else:
                        mask = self.mask
                        pos0min = (numpy.floor(pos[:, :, :, 0].min(axis=-1) - self.offset1).astype(numpy.int32)).clip(0, self._shape_out[0])
                        pos1min = (numpy.floor(pos[:, :, :, 1].min(axis=-1) - self.offset2).astype(numpy.int32)).clip(0, self._shape_out[1])
                        pos0max = (numpy.ceil(pos[:, :, :, 0].max(axis=-1) - self.offset1 + 1).astype(numpy.int32)).clip(0, self._shape_out[0])
                        pos1max = (numpy.ceil(pos[:, :, :, 1].max(axis=-1) - self.offset2 + 1).astype(numpy.int32)).clip(0, self._shape_out[1])
                        self.bin_size = numpy.zeros(self._shape_out, dtype=numpy.int32)
                        for i in range(self.shape_in[0]):
                            for j in range(self.shape_in[1]):
                                if (mask is not None) and mask[i, j]:
                                    continue
                                self.bin_size[pos0min[i, j]:pos0max[i, j], pos1min[i, j]:pos1max[i, j]] += 1
                    self.max_size = self.bin_size.max()
        return self.bin_size

    def calc_init(self):
        """Initialize all arrays
        """
        self.calc_pos()
        self.calc_size()
        self.calc_LUT()
        if ocl and self.device is not None:
            if "lower" in dir(self.device):
                self.device = self.device.lower()
                if self.method == "lut":
                    self.integrator = ocl_azim_lut.OCL_LUT_Integrator(self.lut,
                                                                      self._shape_out[0] * self._shape_out[1],
                                                                      devicetype=self.device)
                else:
                    self.integrator = ocl_azim_csr.OCL_CSR_Integrator(self.lut,
                                                                      self._shape_out[0] * self._shape_out[1],
                                                                      devicetype=self.device,
                                                                      block_size=self.workgroup)
                    self.integrator.workgroup_size["csr_integrate4"] = 1,
            else:
                if self.method == "lut":
                    self.integrator = ocl_azim_lut.OCL_LUT_Integrator(self.lut,
                                                                      self._shape_out[0] * self._shape_out[1],
                                                                      platformid=self.device[0],
                                                                      deviceid=self.device[1])
                else:
                    self.integrator = ocl_azim_csr.OCL_CSR_Integrator(self.lut,
                                                                      self._shape_out[0] * self._shape_out[1],
                                                                      platformid=self.device[0], deviceid=self.device[1],
                                                                      block_size=self.workgroup)
                    self.integrator.workgroup_size["csr_integrate4"] = 1,

    def calc_LUT(self, use_common=True):
        """Calculate the Look-up table

        :return: look up table either in CSR or LUT format depending on serl.method
        """
        if self.pos is None:
            self.calc_pos()

        if self.max_size is None and not use_common:
            self.calc_size()
        if self.lut is None:
            with self._sem:
                if self.lut is None:
                    mask = self.mask
                    if _distortion:
                        if use_common:
                            self.lut = _distortion.calc_sparse(self.pos, self._shape_out, max_pixel_size=(self.delta1, self.delta2), format=self.method)
                        else:
                            if self.method == "lut":
                                self.lut = _distortion.calc_LUT(self.pos, self._shape_out, self.bin_size, max_pixel_size=(self.delta1, self.delta2))
                            else:
                                self.lut = _distortion.calc_CSR(self.pos, self._shape_out, self.bin_size, max_pixel_size=(self.delta1, self.delta2))
                    else:
                        lut = numpy.recarray(shape=(self._shape_out[0], self._shape_out[1], self.max_size), dtype=[("idx", numpy.uint32), ("coef", numpy.float32)])
                        lut[:, :, :].idx = 0
                        lut[:, :, :].coef = 0.0
                        outMax = numpy.zeros(self._shape_out, dtype=numpy.uint32)
                        idx = 0
                        buffer_ = numpy.empty((self.delta1, self.delta2))
                        quad = Quad(buffer_)
                        for i in range(self._shape_out[0]):
                            for j in range(self._shape_out[1]):
                                if (mask is not None) and mask[i, j]:
                                    continue
                                # i,j, idx are indexes of the raw image uncorrected
                                quad.reinit(*list(self.pos[i, j, :, :].ravel()))
                                # print(self.pos[i, j, 0, :], self.pos[i, j, 1, :], self.pos[i, j, 2, :], self.pos[i, j, 3, :]
                                try:
                                    quad.populate_box()
                                except Exception as error:
                                    print("error in quad.populate_box of pixel %i, %i: %s" % (i, j, error))
                                    print("calc_area_vectorial", quad.calc_area_vectorial())
                                    print(self.pos[i, j, 0, :], self.pos[i, j, 1, :], self.pos[i, j, 2, :], self.pos[i, j, 3, :])
                                    print(quad)
                                    raise
                #                box = quad.get_box()
                                for ms in range(quad.get_box_size0()):
                                    ml = ms + quad.get_offset0()
                                    if ml < 0 or ml >= self._shape_out[0]:
                                        continue
                                    for ns in range(quad.get_box_size1()):
                                        # ms,ns are indexes of the corrected image in short form, ml & nl are the same
                                        nl = ns + quad.get_offset1()
                                        if nl < 0 or nl >= self._shape_out[1]:
                                            continue
                                        val = quad.get_box(ms, ns)
                                        if val <= 0:
                                            continue
                                        k = outMax[ml, nl]
                                        lut[ml, nl, k].idx = idx
                                        lut[ml, nl, k].coef = val
                                        outMax[ml, nl] = k + 1
                                idx += 1
                        lut.shape = (self._shape_out[0] * self._shape_out[1]), self.max_size
                        self.lut = lut
        return self.lut

    def correct(self, image, dummy=None, delta_dummy=None):
        """
        Correct an image based on the look-up table calculated ...

        :param image: 2D-array with the image
        :param dummy: value suggested for bad pixels
        :param delta_dummy: precision of the dummy value
        :return: corrected 2D image
        """
        if image.ndim == 2:
            image = resize_image_2D(image, self.shape_in)
        else:  # assume 2d+nchanel
            if _distortion:
                image = _distortion.resize_image_3D(image, self.shape_in)
            else:
                assert image.ndim == 3, "image is 3D"
                shape_in0, shape_in1 = self.shape_in
                shape_img0, shape_img1, nchan = image.shape
                if not ((shape_img0 == shape_in0) and (shape_img1 == shape_in1)):
                    new_image = numpy.zeros((shape_in0, shape_in1, nchan), dtype=numpy.float32)
                    if shape_img0 < shape_in0:
                        if shape_img1 < shape_in1:
                            new_image[:shape_img0, :shape_img1, :] = image
                        else:
                            new_image[:shape_img0, :, :] = image[:, :shape_in1, :]
                    else:
                        if shape_img1 < shape_in1:
                            new_image[:, :shape_img1, :] = image[:shape_in0, :, :]
                        else:
                            new_image[:, :, :] = image[:shape_in0, :shape_in1, :]
                    logger.warning("Patching image of shape %ix%i on expected size of %ix%i",
                                   shape_img1, shape_img0, shape_in1, shape_in0)
                image = new_image
        if self.device:
            if self.integrator is None:
                self.calc_init()
            out = self.integrator.integrate(image)[1]
        else:
            if self.lut is None:
                self.calc_LUT()
            if _distortion is not None:
                out = _distortion.correct(image, self.shape_in, self._shape_out, self.lut,
                                          dummy=dummy or self.empty, delta_dummy=delta_dummy)
            else:
                if self.method == "lut":
                    big = image.ravel().take(self.lut.idx) * self.lut.coef
                    out = big.sum(axis=-1)
                elif self.method == "csr":
                    big = self.lut[0] * image.ravel().take(self.lut[1])
                    indptr = self.lut[2]
                    out = numpy.zeros(indptr.size - 1)
                    for i in range(indptr.size - 1):
                        out[i] = big[indptr[i]:indptr[i + 1]].sum()
        try:
            if image.ndim == 2:
                out.shape = self._shape_out
            else:
                for ds in out:
                    if ds.ndim == 2:
                        ds.shape = self._shape_out
                    else:
                        ds.shape = self._shape_out + ds.shape[2:]

        except ValueError as _err:
            logger.error("Requested in_shape=%s out_shape=%s and ", self.shape_in, self.shape_out)
            raise
        return out

    def correct_ng(self, image,
                   variance=None,
                   dark=None,
                   flat=None,
                   solidangle=None,
                   polarization=None,
                   dummy=None,
                   delta_dummy=None,
                   normalization_factor=1.0):
        """
        Correct an image based on the look-up table calculated ...
        Like the integrate_ng it provides
        * Dark current correction
        * Normalisation with flatfield (or solid angle, polarization, absorption, ...)
        * Error propagation

        :param image: 2D-array with the image
        :param variance: 2D-array with the associated image
        :param dark: array with dark-current values
        :param flat: array with values for a flat image
        :param solidangle: solid-angle array
        :param polarization: numpy array with 2D polarization corrections
        :param dummy: value suggested for bad pixels
        :param delta_dummy: precision of the dummy value
        :param normalization_factor: multiply all normalization with this value
        :return: corrected 2D image
        """
        assert image.ndim == 2
        if variance is not None:
            assert variance.shape == image.shape

        if image.shape != self.shape_in:
            logger.warning("The image shape %s is not the same as the detector %s", image.shape, self.shape_in)
            image = resize_image_2D(image, self.shape_in)
            if variance is not None:
                variance = resize_image_2D(variance, self.shape_in)
            if dark is not None:
                dark = resize_image_2D(dark, self.shape_in)

        if self.device:
            if self.integrator is None:
                self.calc_init()
            res = self.integrator.integrate_ng(image,
                                               variance=variance,
                                               flat=flat,
                                               dark=dark,
                                               solidangle=solidangle,
                                               polarization=polarization,
                                               dummy=dummy,
                                               delta_dummy=delta_dummy,
                                               normalization_factor=normalization_factor,
                                               out_merged=False
                                               )
            if variance is not None:
                if image.ndim == 2:
                    out = res.intensity.reshape(self._shape_out)
                else:
                    out = res.intensity
            else:
                if image.ndim == 2:
                    out = (res.intensity.reshape(self._shape_out), res.error.reshape(self._shape_out))
                else:
                    out = (res.intensity, res.error)
        else:
            if self.lut is None:
                self.calc_LUT()
            if _distortion is not None:
                out = _distortion.correct(image, self.shape_in, self._shape_out, self.lut,
                                          dummy=dummy or self.empty, delta_dummy=delta_dummy)
            else:
                if self.method == "lut":
                    big = image.ravel().take(self.lut.idx) * self.lut.coef
                    out = big.sum(axis=-1)
                elif self.method == "csr":
                    big = self.lut[0] * image.ravel().take(self.lut[1])
                    indptr = self.lut[2]
                    out = numpy.zeros(indptr.size - 1)
                    for i in range(indptr.size - 1):
                        out[i] = big[indptr[i]:indptr[i + 1]].sum()
            try:
                if image.ndim == 2:
                    out.shape = self._shape_out
                else:
                    for ds in out:
                        if ds.ndim == 2:
                            ds.shape = self._shape_out
                        else:
                            ds.shape = self._shape_out + ds.shape[2:]

            except ValueError as _err:
                logger.error("Requested in_shape=%s out_shape=%s and ", self.shape_in, self.shape_out)
                raise
        return out

    def uncorrect(self, image, use_cython=False):
        """
        Take an image which has been corrected and transform it into it's raw (with loss of information)

        :param image: 2D-array with the image
        :return: uncorrected 2D image

        Nota: to retrieve the input mask on can do:

        >>> msk =  dis.uncorrect(numpy.ones(dis._shape_out)) <= 0
        """
        assert image.shape == self._shape_out
        if self.lut is None:
            self.calc_LUT()
        if (linalg is not None) and (use_cython is False):
            if self.method == "lut":
                csr = csr_matrix(sparse_utils.LUT_to_CSR(self.lut))
            else:
                csr = csr_matrix(self.lut)
            res = linalg.lsmr(csr, image.ravel())
            out = res[0].reshape(self.shape_in)
        else:  # This is deprecated and does not work with resise=True
            if self.method == "lut":
                if _distortion is not None:
                    out, _mask = _distortion.uncorrect_LUT(image, self.shape_in, self.lut)
                else:
                    out = numpy.zeros(self.shape_in, dtype=numpy.float32)
                    lout = out.ravel()
                    lin = image.ravel()
                    tot = self.lut.coef.sum(axis=-1)
                    for idx in range(self.lut.shape[0]):
                        t = tot[idx]
                        if t <= 0:
                            continue
                        val = lin[idx] / t
                        lout[self.lut[idx].idx] += val * self.lut[idx].coef
            elif self.method == "csr":
                if _distortion is not None:
                    out, _mask = _distortion.uncorrect_CSR(image, self.shape_in, self.lut)
            else:
                raise NotImplementedError()
        return out

class Quad(object):
    """
    Quad modelisation.

    .. image:: ../img/quad_model.svg
        :alt: Modelization of the quad
    """
    def __init__(self, buffer):
        self.box = buffer
        self.A0 = self.A1 = None
        self.B0 = self.B1 = None
        self.C0 = self.C1 = None
        self.D0 = self.D1 = None
        self.offset0 = self.offset1 = None
        self.box_size0 = self.box_size1 = None
        self.pAB = self.pBC = self.pCD = self.pDA = None
        self.cAB = self.cBC = self.cCD = self.cDA = None
        self.area = None

    def get_idx(self, i, j):
        pass

    def get_box(self, i, j):
        return self.box[i, j]

    def get_offset0(self):
        return self.offset0

    def get_offset1(self):
        return self.offset1

    def get_box_size0(self):
        return self.box_size0

    def get_box_size1(self):
        return self.box_size1

    def reinit(self, A0, A1, B0, B1, C0, C1, D0, D1):
        self.box[:, :] = 0.0
        self.A0 = A0
        self.A1 = A1
        self.B0 = B0
        self.B1 = B1
        self.C0 = C0
        self.C1 = C1
        self.D0 = D0
        self.D1 = D1
        self.offset0 = int(floor(min(self.A0, self.B0, self.C0, self.D0)))
        self.offset1 = int(floor(min(self.A1, self.B1, self.C1, self.D1)))
        self.box_size0 = int(ceil(max(self.A0, self.B0, self.C0, self.D0))) - self.offset0
        self.box_size1 = int(ceil(max(self.A1, self.B1, self.C1, self.D1))) - self.offset1
        self.A0 -= self.offset0
        self.A1 -= self.offset1
        self.B0 -= self.offset0
        self.B1 -= self.offset1
        self.C0 -= self.offset0
        self.C1 -= self.offset1
        self.D0 -= self.offset0
        self.D1 -= self.offset1
        self.pAB = self.pBC = self.pCD = self.pDA = None
        self.cAB = self.cBC = self.cCD = self.cDA = None
        self.area = None

    def __repr__(self):
        return os.linesep.join(["offset %i,%i size %i, %i" % (self.offset0, self.offset1, self.box_size0, self.box_size1), "box: %s" % self.box[:self.box_size0, :self.box_size1]])

    def init_slope(self):
        if self.pAB is None:
            if self.B0 == self.A0:
                self.pAB = numpy.inf
            else:
                self.pAB = (self.B1 - self.A1) / (self.B0 - self.A0)
            if self.C0 == self.B0:
                self.pBC = numpy.inf
            else:
                self.pBC = (self.C1 - self.B1) / (self.C0 - self.B0)
            if self.D0 == self.C0:
                self.pCD = numpy.inf
            else:
                self.pCD = (self.D1 - self.C1) / (self.D0 - self.C0)
            if self.A0 == self.D0:
                self.pDA = numpy.inf
            else:
                self.pDA = (self.A1 - self.D1) / (self.A0 - self.D0)
            self.cAB = self.A1 - self.pAB * self.A0
            self.cBC = self.B1 - self.pBC * self.B0
            self.cCD = self.C1 - self.pCD * self.C0
            self.cDA = self.D1 - self.pDA * self.D0

    def calc_area_AB(self, I1, I2):
        if numpy.isfinite(self.pAB):
            return 0.5 * (I2 - I1) * (self.pAB * (I2 + I1) + 2 * self.cAB)
        else:
            return 0

    def calc_area_BC(self, J1, J2):
        if numpy.isfinite(self.pBC):
            return 0.5 * (J2 - J1) * (self.pBC * (J1 + J2) + 2 * self.cBC)
        else:
            return 0

    def calc_area_CD(self, K1, K2):
        if numpy.isfinite(self.pCD):
            return 0.5 * (K2 - K1) * (self.pCD * (K2 + K1) + 2 * self.cCD)
        else:
            return 0

    def calc_area_DA(self, L1, L2):

        if numpy.isfinite(self.pDA):
            return 0.5 * (L2 - L1) * (self.pDA * (L1 + L2) + 2 * self.cDA)
        else:
            return 0

    def calc_area_old(self):
        if self.area is None:
            if self.pAB is None:
                self.init_slope()
            self.area = -(self.calc_area_AB(self.A0, self.B0) +
                          self.calc_area_BC(self.B0, self.C0) +
                          self.calc_area_CD(self.C0, self.D0) +
                          self.calc_area_DA(self.D0, self.A0))
        return self.area

    def calc_area_vectorial(self):
        if self.area is None:
            self.area = numpy.cross([self.C0 - self.A0, self.C1 - self.A1], [self.D0 - self.B0, self.D1 - self.B1]) / 2.0
        return self.area
    calc_area = calc_area_vectorial

    def populate_box(self):
        if self.pAB is None:
            self.init_slope()
        self.integrateAB(self.B0, self.A0, self.calc_area_AB)
        self.integrateAB(self.A0, self.D0, self.calc_area_DA)
        self.integrateAB(self.D0, self.C0, self.calc_area_CD)
        self.integrateAB(self.C0, self.B0, self.calc_area_BC)
        if (self.box / self.calc_area()).min() < 0:
            print(self.box)
            self.box[:, :] = 0
            print("AB")
            self.integrateAB(self.B0, self.A0, self.calc_area_AB)
            print(self.box)
            self.box[:, :] = 0
            print("DA")
            self.integrateAB(self.A0, self.D0, self.calc_area_DA)
            print(self.box)
            self.box[:, :] = 0
            print("CD")
            self.integrateAB(self.D0, self.C0, self.calc_area_CD)
            print(self.box)
            self.box[:, :] = 0
            print("BC")
            self.integrateAB(self.C0, self.B0, self.calc_area_BC)
            print(self.box)
            print(self)
            raise RuntimeError()
        self.box /= self.calc_area_vectorial()

    def integrateAB(self, start, stop, calc_area):
        h = 0
#        print(start, stop, calc_area(start, stop)
        if start < stop:  # positive contribution
            P = ceil(start)
            dP = P - start
#            print("Integrate", start, P, stop, calc_area(start, stop)
            if P > stop:  # start and stop are in the same unit
                A = calc_area(start, stop)
                if A != 0:
                    AA = abs(A)
                    sign = A / AA
                    dA = (stop - start)  # always positive
#                    print(AA, sign, dA
                    h = 0
                    while AA > 0:
                        if dA > AA:
                            dA = AA
                            AA = -1
                        self.box[int(floor(start)), h] += sign * dA
                        AA -= dA
                        h += 1
            else:
                if dP > 0:
                    A = calc_area(start, P)
                    if A != 0:
                        AA = abs(A)
                        sign = A / AA
                        h = 0
                        dA = dP
                        while AA > 0:
                            if dA > AA:
                                dA = AA
                                AA = -1
                            self.box[int(floor(P)) - 1, h] += sign * dA
                            AA -= dA
                            h += 1
                # subsection P1->Pn
                for i in range(int(floor(P)), int(floor(stop))):
                    A = calc_area(i, i + 1)
                    if A != 0:
                        AA = abs(A)
                        sign = A / AA

                        h = 0
                        dA = 1.0
                        while AA > 0:
                            if dA > AA:
                                dA = AA
                                AA = -1
                            self.box[i, h] += sign * dA
                            AA -= dA
                            h += 1
                # Section Pn->B
                P = floor(stop)
                dP = stop - P
                if dP > 0:
                    A = calc_area(P, stop)
                    if A != 0:
                        AA = abs(A)
                        sign = A / AA
                        h = 0
                        dA = abs(dP)
                        while AA > 0:
                            if dA > AA:
                                dA = AA
                                AA = -1
                            self.box[int(floor(P)), h] += sign * dA
                            AA -= dA
                            h += 1
        elif start > stop:  # negative contribution. Nota is start=stop: no contribution
            P = floor(start)
            if stop > P:  # start and stop are in the same unit
                A = calc_area(start, stop)
                if A != 0:
                    AA = abs(A)
                    sign = A / AA
                    dA = (start - stop)  # always positive
                    h = 0
                    while AA > 0:
                        if dA > AA:
                            dA = AA
                            AA = -1
                        self.box[int(floor(start)), h] += sign * dA
                        AA -= dA
                        h += 1
            else:
                dP = P - start
                if dP < 0:
                    A = calc_area(start, P)
                    if A != 0:
                        AA = abs(A)
                        sign = A / AA
                        h = 0
                        dA = abs(dP)
                        while AA > 0:
                            if dA > AA:
                                dA = AA
                                AA = -1
                            self.box[int(floor(P)), h] += sign * dA
                            AA -= dA
                            h += 1
                # subsection P1->Pn
                for i in range(int(start), int(ceil(stop)), -1):
                    A = calc_area(i, i - 1)
                    if A != 0:
                        AA = abs(A)
                        sign = A / AA
                        h = 0
                        dA = 1
                        while AA > 0:
                            if dA > AA:
                                dA = AA
                                AA = -1
                            self.box[i - 1, h] += sign * dA
                            AA -= dA
                            h += 1
                # Section Pn->B
                P = ceil(stop)
                dP = stop - P
                if dP < 0:
                    A = calc_area(P, stop)
                    if A != 0:
                        AA = abs(A)
                        sign = A / AA
                        h = 0
                        dA = abs(dP)
                        while AA > 0:
                            if dA > AA:
                                dA = AA
                                AA = -1
                            self.box[int(floor(stop)), h] += sign * dA
                            AA -= dA
                            h += 1


#!/usr/bin/env python

from astropy.io import fits as pyfits
import numpy as np
import os
import argparse
from scipy.ndimage import filters

if __name__ == '__main__':

    p   = argparse.ArgumentParser(
        description='''
        Group pha files using optimum binning based on the instrument
        resoluions, or using the formula in Kaastra+Bleeker 16,
        with the addition of a signal to noise constraint.
        The input spectrum is assumed to have the RESPFILE and
        BACKFILE keywords
        ''',
        formatter_class=argparse.ArgumentDefaultsHelpFormatter )

    p.add_argument("spec_file"  , metavar="spec_file", type=str,
            help="The name of the input spectrum file.")
    p.add_argument("out_file"  , metavar="out_file", type=str,
            help="The name of the output spectrum file.")
    p.add_argument('-s', '--snr', metavar='min_snr',
            type=float, default=3,
            help='The minimum signal to noise ration per bin. -1 to ignore')
    p.add_argument('-c', '--counts', metavar='min_counts',
            type=float, default=-1,
            help='The minimum counts per bin. -1 to ignore')
    p.add_argument('-f', '--osample_fac', metavar='osample_fac',
            type=float, default=3,
            help='The oversampling factor in units of the local FWHM')
    p.add_argument("--use_formula", action='store_true', default=False,
            help=('Calculate the oversampling factor using formula 36-37 '
                  'in Kaastra & Bleeker 2016'))
    p.add_argument("--write_chan", action='store_true', default=False,
            help=('Write grouping file, then use standard grppha.'
                  'This requires heasoft and ngroups < 100'))

    # process input #
    args        = p.parse_args()
    spec_file   = args.spec_file
    out_file    = args.out_file
    min_snr     = args.snr
    min_counts  = args.counts
    osample_fac = args.osample_fac
    use_formula = args.use_formula
    write_chan  = args.write_chan

    # some defaults #
    if osample_fac < 0: osample_fac = 1

    # ------------------------------------------- #
    # Read the response and background file names #
    with pyfits.open(spec_file) as fp:
        src_c  = np.array(fp['SPECTRUM'].data.field(1), np.double)
        src_ex = fp['SPECTRUM'].header['EXPOSURE']
        try:
            src_bs = fp['SPECTRUM'].header['BACKSCAL']
        except:
            src_bs = 1.0
        try:
            rsp_file = fp['SPECTRUM'].header['RESPFILE']
        except:
            raise ValueError('No RESPFILE key in spectrum header')
        try:
            bgd_file = fp['SPECTRUM'].header['BACKFILE']
        except:
            bgd_file = None
            #raise ValueError('No BACKFILE key in spectrum header; Set to none if needed')
    # ------------------------------------------- #

    # --------------------- #
    # get background counts #
    bgd_c = np.zeros_like(src_c)
    if not bgd_file in [None, 'none', 'NONE']:
        with pyfits.open(bgd_file) as fp:
            bgd_c  = np.array(fp['SPECTRUM'].data.field(1), np.double)
            bgd_ex = fp['SPECTRUM'].header['EXPOSURE']
            try:
                bgd_bs = fp['SPECTRUM'].header['BACKSCAL']
            except:
                bgd_bs = 1.0
        bgd_c *= (src_bs/bgd_bs) * (src_ex/bgd_ex)
    # --------------------- #

    # ------------------------- #
    # energy-channel conversion #
    with pyfits.open(rsp_file) as fp:
        edata  = fp['EBOUNDS'].data
        chan   = edata.field(0)
        energy = (edata.field(1)+edata.field(2)) / 2.
        try:
            matrix = fp['MATRIX'].data
        except:
            matrix = fp['SPECRESP MATRIX'].data
        nchan  = len(chan)
        nen    = len(matrix)
    men = np.array([(m[1]+m[0])/2 for m in matrix])
    # ------------------------- #

    # --------------------- #
    # loop through channels #
    ibin = np.zeros(nchan) - 1
    ich, ibin[0] = 0, 1
    smooth = nen * 1. / nchan
    while ich < nchan:

        # get the response curve at ich #
        ie = np.argmin(np.abs(men - energy[ich]))
        istart, ilen = matrix[ie][3], matrix[ie][4]
        if not isinstance(istart, (list, np.ndarray)):
            istart, ilen = [istart], [ilen]
        if len(istart) == 0 or len(ilen) == 0:
            ich += 1
            continue
        iarr = np.concatenate([np.arange(i1, i1+i2) for i1,i2 in zip(istart, ilen)])
        rarr = matrix[ie][5]
        rarr = filters.gaussian_filter1d(rarr, smooth)

        # work out fwhm at ich #
        if not np.allclose(rarr, 0) and nchan > 100:
            imax = np.argmax(rarr)

            # limits of fwhm in energy grid units
            ic1 = np.argmin(np.abs(rarr[:imax]-rarr[imax]/2.)) if imax !=0 else 0
            ic2 = np.argmin(np.abs(rarr[imax:]-rarr[imax]/2.)) + imax
            width = ic2 - ic1

        else:
            #width = nchan - ich
            width = 1

        # get oversampling factor if we use_formula is requested #
        if use_formula:
            width = np.max([1, width])
            ind = range(ich, ich+width)
            if ind[-1] >= nchan:
                ind = range(ich, nchan)
            counts = np.max([sum(src_c[ind] - bgd_c[ind]), 1e-10] )
            counts *= 1./width # per resolution element
            x = np.log(counts*(1 + 0.20 *np.log(width)))
            osample_fac = 1. if x<2.119 else (0.08+7./x + 1.8/x**2)/(1+5.9/x)
            osample_fac = 1. / osample_fac
        # ----- #

        # bin width in channel units using oversample_fac
        width = np.int(np.round(np.max([1, width*1./osample_fac])))
        ind   = range(ich, np.min([ich+width, nchan]) )

        # increase width until snr > min_snr #
        if min_snr:
            snr = sum(src_c[ind] - bgd_c[ind])
            while snr < 1 and ich+width<=nchan:
                width += 1
                ind = range(ich, min(ich+width, nchan) )
                snr = sum(src_c[ind] - bgd_c[ind])
            if snr > 0:
                snr /= np.sqrt(sum(src_c[ind] + bgd_c[ind]))
            while (snr < min_snr) and (ich+width < nchan):
                ind = range(ich, ich+width)
                if ind[-1] >= nchan:
                    ind = range(ich, nchan)
                    break

                snr = sum(src_c[ind] - bgd_c[ind])
                if snr <= 0:
                    width += 1; continue
                snr /= np.sqrt(sum(src_c[ind] + bgd_c[ind]))
                width += 1

        # do we have a min_counts requirement? #
        if min_counts:
            counts = sum(src_c[ind] - bgd_c[ind])
            while (counts < min_counts) and (ich+width < nchan):
                ind = range(ich, ich+width)
                if ind[-1] >= nchan:
                    ind = range(ich, nchan)
                    break
                counts = sum(src_c[ind] - bgd_c[ind])
                width += 1
        ich += len(ind)
        if ich < nchan: ibin[ich] = 1
    # --------------------- #

    # ------------------ #
    # write the grouping #

    if write_chan:
        # write an ascii file and use grppha #
        ichan = np.arange(nchan)[ibin==1]
        if ichan[-1] < nchan-1:
            ichan = np.append(ichan, nchan)
        bins = []
        for ich in range(len(ichan)-1):
            bins.append([ichan[ich], ichan[ich+1]-1, ichan[ich+1]-ichan[ich]])
        txt = '\n'.join(['{} {} {}'.format(*x) for x in bins])
        with open('tmp_chans.dat', 'w') as fp: fp.write(txt)
        os.system('rm {0} &> /dev/null'.format(out_file))
        os.system('grppha {} {} "group tmp_chans.dat&exit"'.format(spec_file, out_file))
    else:
        # modify the file with pyfits #
        os.system('rm {} &> /dev/null'.format(out_file))
        os.system('grppha {} {} "group min 20&exit" &> /dev/null'.format(spec_file, out_file))
        with pyfits.open(out_file) as fp:
            hdu = fp['SPECTRUM']
            orig_cols = hdu.columns
            orig_cols['GROUPING'].array = np.array(ibin, np.int)
            cols = pyfits.ColDefs(orig_cols)
            tbl = pyfits.BinTableHDU.from_columns(cols)
            hdu.header.update(tbl.header.copy())
            tbl.header = hdu.header.copy()
            grp = pyfits.HDUList([fp[0],tbl])
            #os.system('rm {0} &> /dev/null'.format(out_file))
            grp.writeto(out_file, overwrite=True)
        print('Grouped file {} written sucessfully'.format(out_file))
    # ------------------ #


import itertools
import logging
import math
import os
import string
import subprocess
import tempfile
from collections import defaultdict, Counter
from dataclasses import dataclass, field
from typing import Dict, Iterable, List, Tuple, Optional, Set

import numpy

from eltetrado.model import Atom3D, Structure3D, Structure2D, BasePair3D, Residue3D, GlycosidicBond, ONZ, \
    GbaTetradClassification, Ion, Direction, LoopType, ONZM, GbaQuadruplexClassification, LoopClassification

logging.basicConfig(level=os.environ.get("LOGLEVEL", "INFO"))

@dataclass(order=True)
class Tetrad:
    @staticmethod
    def is_valid(nt1: Residue3D, nt2: Residue3D, nt3: Residue3D, nt4: Residue3D,
                 pair_dictionary: Dict[Tuple[Residue3D, Residue3D], BasePair3D]) -> bool:
        lw1 = pair_dictionary[(nt1, nt2)].lw
        lw2 = pair_dictionary[(nt2, nt3)].lw
        lw3 = pair_dictionary[(nt3, nt4)].lw
        lw4 = pair_dictionary[(nt4, nt1)].lw
        for lw_i, lw_j in ((lw1, lw4), (lw2, lw1), (lw3, lw2), (lw4, lw3)):
            if lw_i.name[1] == lw_j.name[2]:
                return False
        return True

    nt1: Residue3D
    nt2: Residue3D
    nt3: Residue3D
    nt4: Residue3D
    pair_12: BasePair3D
    pair_23: BasePair3D
    pair_34: BasePair3D
    pair_41: BasePair3D
    onz: ONZ = field(init=False)
    gba_class: Optional[GbaTetradClassification] = field(init=False)
    planarity_deviation: float = field(init=False)
    ions_channel: List[Atom3D] = field(default_factory=list)
    ions_outside: Dict[Residue3D, List[Atom3D]] = field(default_factory=dict)

    def __post_init__(self):
        self.reorder_to_match_5p_3p()
        self.planarity_deviation = self.__calculate_planarity_deviation()

    def reorder_to_match_5p_3p(self):
        # transform into (0, 1, 2, 3)
        ni, nj, nk, nl = map(lambda nt: nt.index, self.nucleotides)
        indices = sorted((ni, nj, nk, nl))
        ni, nj, nk, nl = (indices.index(x) for x in (ni, nj, nk, nl))

        nmin = min(ni, nj, nk, nl)
        if nmin == ni:
            pass
        elif nmin == nj:
            self.nt1, self.nt2, self.nt3, self.nt4 = self.nt2, self.nt3, self.nt4, self.nt1
            self.pair_12, self.pair_23, self.pair_34, self.pair_41 = self.pair_23, self.pair_34, self.pair_41, self.pair_12
        elif nmin == nk:
            self.nt1, self.nt2, self.nt3, self.nt4 = self.nt3, self.nt4, self.nt1, self.nt2
            self.pair_12, self.pair_23, self.pair_34, self.pair_41 = self.pair_34, self.pair_41, self.pair_12, self.pair_23
        else:
            self.nt1, self.nt2, self.nt3, self.nt4 = self.nt4, self.nt1, self.nt2, self.nt3
            self.pair_12, self.pair_23, self.pair_34, self.pair_41 = self.pair_41, self.pair_12, self.pair_23, self.pair_34

        # flip order if necessary
        if self.pair_12.score() > self.pair_41.reverse().score():
            self.nt1, self.nt2, self.nt3, self.nt4 = self.nt1, self.nt4, self.nt3, self.nt2
            self.pair_12, self.pair_23, self.pair_34, self.pair_41 = self.pair_41.reverse(), self.pair_34.reverse(), self.pair_23.reverse(), self.pair_12.reverse()

        # ONZ and da Silva's classification are valid in 5'-3' order
        self.onz = self.__classify_onz()
        self.gba_class = self.__classify_by_gba()

    def reorder_to_match_other_tetrad(self, order: Tuple[Residue3D, Residue3D, Residue3D, Residue3D]):
        if order == (self.nt1, self.nt2, self.nt3, self.nt4):
            pass
        elif order == (self.nt2, self.nt3, self.nt4, self.nt1):
            self.pair_12, self.pair_23, self.pair_34, self.pair_41 = self.pair_23, self.pair_34, self.pair_41, self.pair_12
        elif order == (self.nt3, self.nt4, self.nt1, self.nt2):
            self.pair_12, self.pair_23, self.pair_34, self.pair_41 = self.pair_34, self.pair_41, self.pair_12, self.pair_23
        elif order == (self.nt4, self.nt1, self.nt2, self.nt3):
            self.pair_12, self.pair_23, self.pair_34, self.pair_41 = self.pair_41, self.pair_12, self.pair_23, self.pair_34
        elif order == (self.nt4, self.nt3, self.nt2, self.nt1):
            self.pair_12, self.pair_23, self.pair_34, self.pair_41 = self.pair_34.reverse(), self.pair_23.reverse(), self.pair_12.reverse(), self.pair_41.reverse()
        elif order == (self.nt3, self.nt2, self.nt1, self.nt4):
            self.pair_12, self.pair_23, self.pair_34, self.pair_41 = self.pair_23.reverse(), self.pair_12.reverse(), self.pair_41.reverse(), self.pair_34.reverse()
        elif order == (self.nt2, self.nt1, self.nt4, self.nt3):
            self.pair_12, self.pair_23, self.pair_34, self.pair_41 = self.pair_12.reverse(), self.pair_41.reverse(), self.pair_34.reverse(), self.pair_23.reverse()
        elif order == (self.nt1, self.nt4, self.nt3, self.nt2):
            self.pair_12, self.pair_23, self.pair_34, self.pair_41 = self.pair_41.reverse(), self.pair_34.reverse(), self.pair_23.reverse(), self.pair_12.reverse()
        else:
            raise RuntimeError(f'Cannot apply order: {order}')

        self.nt1, self.nt2, self.nt3, self.nt4 = order

    def __classify_onz(self) -> ONZ:
        # transform into (0, 1, 2, 3)
        ni, nj, nk, nl = (nt.index for nt in self.nucleotides)
        indices = sorted((ni, nj, nk, nl))
        ni, nj, nk, nl = (indices.index(x) for x in (ni, nj, nk, nl))

        while ni != 0:
            ni, nj, nk, nl = nl, ni, nj, nk

        order = (nj, nk, nl)
        if order == (1, 2, 3):
            return ONZ.O_PLUS
        elif order == (3, 2, 1):
            return ONZ.O_MINUS
        elif order == (1, 3, 2):
            return ONZ.N_PLUS
        elif order == (2, 3, 1):
            return ONZ.N_MINUS
        elif order == (2, 1, 3):
            return ONZ.Z_PLUS
        elif order == (3, 1, 2):
            return ONZ.Z_MINUS

        raise RuntimeError(f'Impossible combination: {ni} {nj} {nk} {nl}')

    def __classify_by_gba(self) -> Optional[GbaTetradClassification]:
        """
        See: Webba da Silva, M. (2007). Geometric Formalism for DNA Quadruplex Folding.
        Chemistry - A European Journal, 13(35), 9738–9745. https://doi.org/10.1002/chem.200701255

        :return: Classification according to Webba da Silva or n/a
        """
        # without all nucleotides having a valid syn/anti, this classification is impossible
        if not all([nt.chi_class in (GlycosidicBond.syn, GlycosidicBond.anti) for nt in self.nucleotides]):
            return None

        # this will create a 4-letter string made of 's' for syn or 'a' for anti
        fingerprint = ''.join([nt.chi_class.value[0] for nt in self.nucleotides])

        # this dict has all classes mapped to fingerprints
        gba_classes = {
            'aass': GbaTetradClassification.Ia,
            'ssaa': GbaTetradClassification.Ib,
            'asas': GbaTetradClassification.IIa,
            'sasa': GbaTetradClassification.IIb,
            'asaa': GbaTetradClassification.IIIa,
            'sass': GbaTetradClassification.IIIb,
            'aaas': GbaTetradClassification.IVa,
            'sssa': GbaTetradClassification.IVb,
            'aasa': GbaTetradClassification.Va,
            'ssas': GbaTetradClassification.Vb,
            'assa': GbaTetradClassification.VIa,
            'saas': GbaTetradClassification.VIb,
            'asss': GbaTetradClassification.VIIa,
            'saaa': GbaTetradClassification.VIIb,
            'aaaa': GbaTetradClassification.VIIIa,
            'ssss': GbaTetradClassification.VIIIb
        }

        if fingerprint not in gba_classes:
            logging.error(f'Impossible combination of syn/anti: {[nt.chi_class for nt in self.nucleotides]}')
            return None
        return gba_classes[fingerprint]

    def __calculate_planarity_deviation(self) -> float:
        outer = [nt.outermost_atom for nt in self.nucleotides]
        inner = [nt.innermost_atom for nt in self.nucleotides]
        return numpy.linalg.norm(center_of_mass(outer) - center_of_mass(inner))

    @property
    def nucleotides(self) -> Tuple[Residue3D, Residue3D, Residue3D, Residue3D]:
        return self.nt1, self.nt2, self.nt3, self.nt4

    def __hash__(self):
        return hash(frozenset([self.nt1, self.nt2, self.nt3, self.nt4]))

    def __str__(self):
        return f'    ' \
               f'{self.nt1.full_name} {self.nt2.full_name} {self.nt3.full_name} {self.nt4.full_name} ' \
               f'{self.pair_12.lw.value} {self.pair_23.lw.value} {self.pair_34.lw.value} {self.pair_41.lw.value} ' \
               f'{self.onz.value} {self.gba_class.value} ' \
               f'planarity={round(self.planarity_deviation, 2)} ' \
               f'{self.__ions_channel_str()} ' \
               f'{self.__ions_outside_str()}\n'

    def chains(self) -> Set[str]:
        return set([nt.chain for nt in self.nucleotides])

    def is_disjoint(self, other) -> bool:
        return frozenset(self.nucleotides).isdisjoint(frozenset(other.nucleotides))

    def center(self) -> numpy.ndarray:
        return center_of_mass(self.outer_and_inner_atoms())

    def outer_and_inner_atoms(self) -> List[Atom3D]:
        return list(map(lambda residue: residue.outermost_atom, self.nucleotides)) + \
               list(map(lambda residue: residue.innermost_atom, self.nucleotides))

    def __ions_channel_str(self) -> str:
        if self.ions_channel:
            return 'ions_channel=' + ','.join([atom.atomName for atom in self.ions_channel])
        return ''

    def __ions_outside_str(self) -> str:
        if self.ions_outside:
            result = []
            for residue, ions in self.ions_outside.items():
                result.append(f'{residue.full_name}: [{",".join([ion.atomName for ion in ions])}]')
            return 'ions_outside=' + ' '.join(result)
        return ''

@dataclass
class TetradPair:
    tetrad1: Tetrad
    tetrad2: Tetrad
    stacked: Dict[Residue3D, Residue3D]
    tetrad2_nts_best_order: Tuple[Residue3D, Residue3D, Residue3D, Residue3D] = field(init=False)
    direction: Direction = field(init=False)
    rise: float = field(init=False)
    twist: float = field(init=False)

    def __post_init__(self):
        self.tetrad2_nts_best_order = (
            self.stacked[self.tetrad1.nt1], self.stacked[self.tetrad1.nt2],
            self.stacked[self.tetrad1.nt3], self.stacked[self.tetrad1.nt4]
        )
        self.direction = self.__determine_direction()
        self.rise = self.__calculate_rise()
        self.twist = self.__calculate_twist()

    def __determine_direction(self) -> Direction:
        indices1 = list(map(lambda nt: nt.index, self.tetrad1.nucleotides))
        indices2 = list(map(lambda nt: nt.index, self.tetrad2_nts_best_order))

        # count directions 5' -> 3' as +1 or -1
        counter = Counter(1 if j - i > 0 else -1 for i, j in zip(indices1, indices2))
        direction, count = counter.most_common()[0]

        if count == 4:
            # all in the same direction
            return Direction.parallel
        elif count == 2:
            # two in +, one in - direction
            return Direction.antiparallel

        return Direction.hybrid

    def __calculate_rise(self) -> float:
        t1 = self.tetrad1.outer_and_inner_atoms()
        t2 = self.tetrad2.outer_and_inner_atoms()
        return numpy.linalg.norm(center_of_mass(t1) - center_of_mass(t2))

    def __calculate_twist(self) -> float:
        nt1_1, nt1_2, _, _ = self.tetrad1.nucleotides
        nt2_1, nt2_2, _, _ = self.tetrad2_nts_best_order

        v1 = nt1_1.find_atom("C1'").coordinates() - nt1_2.find_atom("C1'").coordinates()
        v1 = v1 / numpy.linalg.norm(v1)
        v2 = nt2_1.find_atom("C1'").coordinates() - nt2_2.find_atom("C1'").coordinates()
        v2 = v2 / numpy.linalg.norm(v2)
        return math.degrees(numpy.arccos(numpy.clip(numpy.dot(v1, v2), -1.0, 1.0)))

    def __str__(self):
        return f'      direction={self.direction.value} rise={round(self.rise, 2)} twist={round(self.twist, 2)}\n'

@dataclass
class Tract:
    nucleotides: List[Residue3D]

    def __str__(self):
        return f'      {", ".join(map(lambda nt: nt.full_name, self.nucleotides))}'

@dataclass
class Loop:
    nucleotides: List[Residue3D]
    loop_type: Optional[LoopType]

    def __str__(self):
        return f'      {self.loop_type.value if self.loop_type else "n/a"} ' \
               f'{", ".join(map(lambda nt: nt.full_name, self.nucleotides))}'

@dataclass
class Quadruplex:
    tetrads: List[Tetrad]
    tetrad_pairs: List[TetradPair]
    structure3d: Structure3D
    onzm: Optional[ONZM] = field(init=False)
    gba_classes: List[GbaQuadruplexClassification] = field(init=False)
    tracts: List[Tract] = field(init=False)
    loops: List[Loop] = field(init=False)
    loop_class: Optional[LoopClassification] = field(init=False)

    def __post_init__(self):
        self.onzm = self.__classify_onzm()
        self.gba_classes = self.__classify_by_gba()
        self.tracts = self.__find_tracts()
        self.loops = self.__find_loops()
        self.loop_class = self.__classify_by_loops()

    def __classify_onzm(self) -> Optional[ONZM]:
        if len(self.tetrads) == 1:
            return None
        if any([t.onz is None for t in self.tetrads]):
            return None

        counter = Counter([t.onz.value[0] for t in self.tetrads])
        onz, support = counter.most_common()[0]
        if support != len(self.tetrads):
            onz = 'M'

        counter = Counter([tp.direction.value[0] for tp in self.tetrad_pairs])
        direction, support = counter.most_common()[0]
        if support != len(self.tetrad_pairs):
            direction = 'h'

        counter = Counter([t.onz.value[1] for t in self.tetrads])
        plus_minus, support = counter.most_common()[0]
        if support != len(self.tetrads):
            plus_minus = '*'

        return ONZM.from_value(f'{onz}{direction}{plus_minus}')

    def __classify_by_gba(self) -> List[GbaQuadruplexClassification]:
        gbas = set()
        for t in self.tetrads:
            gba = t.gba_class
            if gba is not None:
                gbas.add(gba.value[:-1])  # discard 'a' or 'b' subvariant
        roman_numerals = {'I': 1, 'II': 2, 'III': 3, 'IV': 4, 'V': 5, 'VI': 6, 'VII': 7, 'VIII': 8}
        gbas = sorted(gbas, key=lambda gba: roman_numerals.get(gba, 100))
        return list(map(lambda x: GbaQuadruplexClassification[x], gbas))

    def __find_tracts(self) -> List[Tract]:
        tracts = [[self.tetrads[0].nt1], [self.tetrads[0].nt2], [self.tetrads[0].nt3], [self.tetrads[0].nt4]]
        if len(self.tetrad_pairs) > 0:
            for tetrad_pair in self.tetrad_pairs:
                nt_dict = {
                    tetrad_pair.tetrad1.nt1: tetrad_pair.tetrad2_nts_best_order[0],
                    tetrad_pair.tetrad1.nt2: tetrad_pair.tetrad2_nts_best_order[1],
                    tetrad_pair.tetrad1.nt3: tetrad_pair.tetrad2_nts_best_order[2],
                    tetrad_pair.tetrad1.nt4: tetrad_pair.tetrad2_nts_best_order[3],
                }
                for i in range(4):
                    tracts[i].append(nt_dict[tracts[i][-1]])
        return [Tract(nts) for nts in tracts]

    def __find_loops(self) -> List[Loop]:
        if len(self.tetrads) == 1:
            return []

        loops = []
        tetrad_nucleotides = sorted([nt for tetrad in self.tetrads for nt in tetrad.nucleotides],
                                    key=lambda nt: nt.index)

        for i in range(1, len(tetrad_nucleotides)):
            nprev = tetrad_nucleotides[i - 1]
            ncur = tetrad_nucleotides[i]
            if ncur.index - nprev.index > 1 and ncur.chain == nprev.chain:
                for tract in self.tracts:
                    if nprev in tract.nucleotides and ncur in tract.nucleotides:
                        break
                else:
                    nts = list(filter(lambda nt: nprev.index < nt.index < ncur.index, self.structure3d.residues))
                    loop_type = self.__detect_loop_type(nprev, ncur)
                    loops.append(Loop(nts, loop_type))
        return loops

    def __detect_loop_type(self, nt_first: Residue3D, nt_last: Residue3D) -> Optional[LoopType]:
        tetrad_with_first = self.__find_tetrad_with_nt(nt_first)
        tetrad_with_last = self.__find_tetrad_with_nt(nt_last)

        if tetrad_with_first is None or tetrad_with_last is None:
            logging.warning(f'Failed to classify the loop between {nt_first} and {nt_last}')
            return None

        if tetrad_with_first == tetrad_with_last:
            # diagonal or laterals happen when first and last nt of a loop is in the same tetrad
            sign = self.__detect_loop_sign(nt_first, nt_last, tetrad_with_first)
            if sign is not None:
                return LoopType.from_value(f'lateral{sign}')
            return LoopType.diagonal

        tract_with_last = self.__find_tract_with_nt(nt_last)
        if tract_with_last is not None:
            # search along the tract to check what pairs with nt_first
            for nt in tract_with_last.nucleotides:
                if nt in tetrad_with_first.nucleotides:
                    sign = self.__detect_loop_sign(nt_first, nt, tetrad_with_first)
                    if sign is not None:
                        return LoopType.from_value(f'propeller{sign}')
        logging.warning(f'Failed to classify the loop between {nt_first} and {nt_last}')
        return None

    def __find_tetrad_with_nt(self, nt: Residue3D) -> Optional[Tetrad]:
        for tetrad in self.tetrads:
            if nt in tetrad.nucleotides:
                return tetrad
        return None

    def __find_tract_with_nt(self, nt: Residue3D) -> Optional[Tract]:
        for tract in self.tracts:
            if nt in tract.nucleotides:
                return tract
        return None

    def __detect_loop_sign(self, first: Residue3D, last: Residue3D, tetrad: Tetrad) -> Optional[str]:
        for pair in [tetrad.pair_12, tetrad.pair_23, tetrad.pair_34, tetrad.pair_41]:
            # main check
            if pair.nt1 == first and pair.nt2 == last:
                if pair.score() < pair.reverse().score():
                    return '-'
                return '+'
            # reverse check
            if pair.nt1 == last and pair.nt2 == first:
                if pair.score() < pair.reverse().score():
                    return '+'
                return '-'
        return None

    def __classify_by_loops(self) -> Optional[LoopClassification]:
        if len(self.loops) != 3 or any([loop.loop_type is None for loop in self.loops]):
            return None

        loop_classes = {
            'ppp': '1',
            'ppl': '2',
            'plp': '3',
            'lpp': '4',
            'pdp': '5',
            'lll': '6',
            'llp': '7',
            'lpl': '8',
            'pll': '9',
            'pdl': '10',
            'ldl': '11',
            'dpd': '12',
            'ldp': '13'
        }
        fingerprint = ''.join([loop.loop_type.value[0] for loop in self.loops])
        if fingerprint not in loop_classes:
            logging.error(f'Unknown loop classification: {fingerprint}')
            return None
        subtype = 'a' if self.loops[0 if fingerprint != 'dpd' else 1].loop_type.value[-1] == '-' else 'b'
        return LoopClassification.from_value(f'{loop_classes[fingerprint]}{subtype}')

    def __str__(self):
        builder = ''
        if len(self.tetrads) == 1:
            builder += '  single tetrad\n'
            builder += str(self.tetrads[0])
        else:
            builder += f'  {self.onzm.value if self.onzm is not None else "R"}'
            builder += f' {",".join(map(lambda gba: gba.value, self.gba_classes))}'
            if self.loop_class:
                builder += f' {self.loop_class.value} {self.loop_class.loop_progression()}'
            else:
                builder += f' n/a'
            builder += f' quadruplex with {len(self.tetrads)} tetrads\n'
            builder += str(self.tetrad_pairs[0].tetrad1)
            for tetrad_pair in self.tetrad_pairs:
                builder += str(tetrad_pair)
                builder += str(tetrad_pair.tetrad2)
            if self.tracts:
                builder += '\n    Tracts:\n'
                for tract in self.tracts:
                    builder += f'{tract}\n'
            if self.loops:
                builder += '\n    Loops:\n'
                for loop in self.loops:
                    builder += f'{loop}\n'
            builder += '\n'
        return builder

@dataclass
class Helix:
    tetrads: List[Tetrad]
    tetrad_pairs: List[TetradPair]
    structure3d: Structure3D
    quadruplexes: List[Quadruplex] = field(init=False)

    def __post_init__(self):
        self.quadruplexes = self.__find_quadruplexes()

    def __find_quadruplexes(self):
        if len(self.tetrad_pairs) == 0:
            return [Quadruplex(self.tetrads, [], self.structure3d)]

        quadruplexes = list()
        tetrads = list()

        for tetrad in [self.tetrad_pairs[0].tetrad1] + [tetrad_pair.tetrad2 for tetrad_pair in self.tetrad_pairs]:
            if tetrads:
                if tetrad.chains().isdisjoint(tetrads[-1].chains()):
                    quadruplexes.append(Quadruplex(tetrads, self.__filter_tetrad_pairs(tetrads), self.structure3d))
                    tetrads = list()
            tetrads.append(tetrad)

        quadruplexes.append(Quadruplex(tetrads, self.__filter_tetrad_pairs(tetrads), self.structure3d))

        return quadruplexes

    def __filter_tetrad_pairs(self, tetrads: List[Tetrad]) -> List[TetradPair]:
        chains = set()
        for tetrad in tetrads:
            chains.update(tetrad.chains())

        def check_tetrad(t: Tetrad) -> bool:
            return not t.chains().isdisjoint(chains)

        def check_pair(tp: TetradPair) -> bool:
            return check_tetrad(tp.tetrad1) and check_tetrad(tp.tetrad2)

        return list(filter(check_pair, self.tetrad_pairs))

    def __str__(self):
        builder = ''
        if len(self.tetrads) > 1:
            builder += f'n4-helix with {len(self.tetrads)} tetrads\n'
            for quadruplex in self.quadruplexes:
                builder += str(quadruplex)
        elif len(self.tetrads) == 1:
            builder += 'single tetrad without stacking\n'
            builder += str(self.tetrads[0])
        return builder

@dataclass
class Analysis:
    structure2d: Structure2D
    structure3d: Structure3D
    strict: bool
    no_reorder: bool
    stacking_mismatch: int
    base_pairs: List[BasePair3D] = field(init=False)
    base_pair_graph: Dict[Residue3D, List[Residue3D]] = field(init=False)
    base_pair_dict: Dict[Tuple[Residue3D, Residue3D], BasePair3D] = field(init=False)
    stacking_graph: Dict[Residue3D, List[Residue3D]] = field(init=False)
    tetrads: List[Tetrad] = field(init=False)
    tetrad_scores: Dict[Tetrad, Dict[Tetrad, Tuple[int, Tuple, Tuple]]] = field(init=False)
    tetrad_pairs: List[TetradPair] = field(init=False)
    helices: List[Helix] = field(init=False)
    ions: List[Atom3D] = field(init=False)
    sequence: str = field(init=False)
    line1: str = field(init=False)
    line2: str = field(init=False)
    shifts: Dict[Residue3D, int] = field(init=False)

    def __post_init__(self):
        self.base_pairs = self.structure3d.base_pairs(self.structure2d)
        self.base_pair_graph = self.structure3d.base_pair_graph(self.structure2d, self.strict)
        self.base_pair_dict = self.structure3d.base_pair_dict(self.structure2d, self.strict)
        self.stacking_graph = self.structure3d.stacking_graph(self.structure2d)
        self.tetrads = self.__find_tetrads(self.no_reorder)
        self.tetrad_scores = self.__calculate_tetrad_scores()
        self.tetrad_pairs = self.__find_tetrad_pairs(self.stacking_mismatch)
        self.helices = self.__find_helices()

        if not self.no_reorder:
            self.__find_best_chain_order()

        self.sequence, self.line1, self.line2, self.shifts = self.__generate_twoline_dotbracket()
        self.ions = self.__find_ions()
        self.__assign_ions_to_tetrads()

    def __find_tetrads(self, no_reorder=False) -> List[Tetrad]:
        # search for a tetrad: i -> j -> k -> l
        #                      ^--------------^
        tetrads = []
        for i in self.base_pair_graph:
            for j in filter(lambda x: x != i, self.base_pair_graph[i]):
                for k in filter(lambda x: x not in (i, j), self.base_pair_graph[j]):
                    for l in filter(lambda x: x not in (i, j, k) and i in self.base_pair_graph[x],
                                    self.base_pair_graph[k]):
                        if Tetrad.is_valid(i, j, k, l, self.base_pair_dict):
                            pair_12 = self.base_pair_dict[(i, j)]
                            pair_23 = self.base_pair_dict[(j, k)]
                            pair_34 = self.base_pair_dict[(k, l)]
                            pair_41 = self.base_pair_dict[(l, i)]
                            tetrads.append(Tetrad(i, j, k, l, pair_12, pair_23, pair_34, pair_41))

        # build graph of tetrads
        while tetrads:
            graph = defaultdict(list)
            for (ti, tj) in itertools.combinations(tetrads, 2):
                if not ti.is_disjoint(tj):
                    graph[ti].append(tj)
                    graph[tj].append(ti)

            # remove tetrad which conflicts the most with others
            # in case of a tie, remove one which has the worst planarity deviation
            candidates = sorted(tetrads, key=lambda t: (len(graph[t]), t.planarity_deviation),
                                reverse=True)
            if len(graph[candidates[0]]) > 0:
                tetrads.remove(candidates[0])
            else:
                break

        return sorted(tetrads, key=lambda t: min(map(lambda nt: nt.index, t.nucleotides)))

    def __calculate_tetrad_scores(self) \
            -> Dict[Tetrad, Dict[Tetrad, Tuple[int, Tuple, Tuple]]]:
        def is_next_by_stacking(nt1: Residue3D, nt2: Residue3D) -> bool:
            return nt2 in self.stacking_graph.get(nt1, [])

        def is_next_sequentially(nt1: Residue3D, nt2: Residue3D) -> bool:
            return nt1.chain == nt2.chain and abs(nt1.index - nt2.index) == 1

        tetrad_scores = defaultdict(dict)

        for ti, tj in itertools.combinations(self.tetrads, 2):
            nts1 = ti.nucleotides
            best_score = 0
            best_score_sequential = 0
            best_score_stacking = 0
            best_order = tj.nucleotides

            n1, n2, n3, n4 = tj.nucleotides
            viable_permutations = [(n1, n2, n3, n4), (n2, n3, n4, n1), (n3, n4, n1, n2), (n4, n1, n2, n3),
                                   (n1, n4, n3, n2), (n4, n3, n2, n1), (n3, n2, n1, n4), (n2, n1, n4, n3)]

            for nts2 in viable_permutations:
                score_stacking = [1 if is_next_by_stacking(nts1[i], nts2[i]) else 0 for i in range(4)]
                score_sequential = [1 if is_next_sequentially(nts1[i], nts2[i]) else 0 for i in range(4)]
                score = sum([max(score_stacking[i], score_sequential[i]) for i in range(4)])
                score_sequential = sum(score_sequential)
                score_stacking = sum(score_stacking)

                if (score, score_sequential, score_stacking) > (best_score, best_score_sequential, best_score_stacking):
                    best_score, best_score_sequential, best_score_stacking = score, score_sequential, score_stacking
                    best_order = nts2
                if best_score == 4:
                    break

            tetrad_scores[ti][tj] = (best_score, nts1, best_order)
            tetrad_scores[tj][ti] = (best_score, best_order, nts1)

        return tetrad_scores

    def __find_tetrad_pairs(self, stacking_mismatch: int) -> List[TetradPair]:
        tetrads = list(self.tetrads)
        best_score = 0
        best_order = tetrads

        for ti in tetrads:
            score = 0
            order = [ti]
            candidates = set(self.tetrads) - {ti}

            while candidates:
                tj = max([tj for tj in candidates], key=lambda tk: self.tetrad_scores[ti][tk][0])
                score += self.tetrad_scores[ti][tj][0]
                order.append(tj)
                candidates.remove(tj)
                ti = tj

            if score > best_score:
                best_score = score
                best_order = order

            if best_score == (len(self.tetrads) - 1) * 4:
                break

        tetrad_pairs = []

        for i in range(1, len(best_order)):
            ti, tj = best_order[i - 1], best_order[i]
            score = self.tetrad_scores[ti][tj][0]

            if score >= (4 - stacking_mismatch):
                nts1, nts2 = self.tetrad_scores[ti][tj][1:]
                stacked = {nts1[i]: nts2[i] for i in range(4)}
                stacked.update({v: k for k, v in stacked.items()})
                tetrad_pairs.append(TetradPair(ti, tj, stacked))
                order = (stacked[ti.nt1], stacked[ti.nt2], stacked[ti.nt3], stacked[ti.nt4])
                tj.reorder_to_match_other_tetrad(order)

        return tetrad_pairs

    def __find_helices(self):
        helices = []
        helix_tetrads = []
        helix_tetrad_pairs = []

        for tp in self.tetrad_pairs:
            ti, tj = tp.tetrad1, tp.tetrad2
            if not helix_tetrads:
                helix_tetrads.append(ti)
            score = self.tetrad_scores[helix_tetrads[-1]][tj][0]
            if score >= (4 - self.stacking_mismatch):
                helix_tetrads.append(tj)
                helix_tetrad_pairs.append(tp)
            else:
                helices.append(Helix(helix_tetrads, helix_tetrad_pairs, self.structure3d))
                helix_tetrads = [ti, tj]
                helix_tetrad_pairs = [tp]

        if helix_tetrads:
            helices.append(Helix(helix_tetrads, helix_tetrad_pairs, self.structure3d))

        for tetrad in self.tetrads:
            if not any([tetrad in helix.tetrads for helix in helices]):
                helices.append(Helix([tetrad], [], self.structure3d))

        return helices

    def __find_best_chain_order(self):
        chain_groups = self.__group_related_chains()
        final_order = []

        for chains in chain_groups:
            best_permutation, best_score = chains, (1e10, 1e10)

            if len(chains) > 1:
                for permutation in itertools.permutations(chains):
                    self.__reorder_chains(permutation)
                    classifications = [t.onz for h in self.helices for t in h.tetrads]
                    logging.debug(
                        f'Checking reorder: {" ".join(permutation)} {" ".join(map(lambda c: c.value, classifications))}')

                    onz_score = sum(c.score() for c in classifications)
                    chain_order_score = self.__chain_order_score(permutation)
                    score = (onz_score, chain_order_score)

                    if score < best_score:
                        best_score = score
                        best_permutation = permutation
                    elif score == best_score:
                        # in case of a tie, pick permutation earlier in lexicographical sense
                        if permutation < best_permutation:
                            best_permutation = permutation

            final_order.extend(best_permutation)

        if len(final_order) > 1:
            self.__reorder_chains(final_order)
            classifications = [t.onz for h in self.helices for t in h.tetrads]
            logging.debug(f'Selected chain order: {" ".join(final_order)} '
                          f'{" ".join(map(lambda onz: onz.value, classifications))}')

            self.tetrads = self.__find_tetrads(True)
            self.tetrad_scores = self.__calculate_tetrad_scores()
            self.tetrad_pairs = self.__find_tetrad_pairs(self.stacking_mismatch)
            self.helices = self.__find_helices()

    def __group_related_chains(self) -> List[List[str]]:
        candidates = set()

        for h in self.helices:
            for t in h.tetrads:
                candidates.add(frozenset([t.nt1.chain, t.nt2.chain, t.nt3.chain, t.nt4.chain]))

        candidates = [set(c) for c in candidates]
        changed = True

        while changed:
            changed = False

            for i, j in itertools.combinations(range(len(candidates)), 2):
                qi, qj = candidates[i], candidates[j]

                if not qi.isdisjoint(qj):
                    qi.update(qj)
                    del candidates[j]
                    changed = True
                    break

        candidates = sorted(candidates, key=lambda x: len(x), reverse=True)
        groups = []

        for candidate in candidates:
            if any([group.issuperset(candidate) for group in groups]):
                continue
            groups.append(candidate)

        return sorted([sorted(group) for group in groups], key=lambda x: x[0])

    def __reorder_chains(self, chain_order: Iterable[str]):
        i = 1
        for chain in chain_order:
            for nt in self.structure3d.residues:
                if nt.chain == chain:
                    nt.index = i
                    i += 1
        for nt in self.structure3d.residues:
            if nt.chain not in chain_order:
                nt.index = i
                i += 1

        if len(self.tetrad_pairs) > 0:
            self.tetrad_pairs[0].tetrad1.reorder_to_match_5p_3p()
            for tp in self.tetrad_pairs:
                order = (tp.stacked[tp.tetrad1.nt1], tp.stacked[tp.tetrad1.nt2],
                         tp.stacked[tp.tetrad1.nt3], tp.stacked[tp.tetrad1.nt4])
                tp.tetrad2.reorder_to_match_5p_3p()  # this is required to recalculate ONZ
                tp.tetrad2.reorder_to_match_other_tetrad(order)

    def __chain_order_score(self, chain_order: Tuple[str, ...]) -> int:
        chain_pairs = []
        for h in self.helices:
            for t in h.tetrads:
                for p in [t.pair_12, t.pair_23, t.pair_34, t.pair_41]:
                    c1 = p.nt1.chain
                    c2 = p.nt2.chain
                    if c1 != c2 and c1 in chain_order and c2 in chain_order:
                        chain_pairs.append([c1, c2])
        sum_sq = 0
        for c1, c2 in chain_pairs:
            sum_sq += (chain_order.index(c1) - chain_order.index(c2)) ** 2
        return sum_sq

    def __find_ions(self) -> List[Atom3D]:
        metal_atom_names = set([ion.value.upper() for ion in Ion])
        ions = []
        used = set()
        for residue in self.structure3d.residues:
            for atom in residue.atoms:
                if atom.atomName.upper() in metal_atom_names:
                    coordinates = tuple(atom.coordinates())
                    if coordinates not in used:
                        ions.append(atom)
                        used.add(coordinates)
        return ions

    def __assign_ions_to_tetrads(self) \
            -> Tuple[Dict[Tetrad, List[Atom3D]], Dict[Tuple[Tetrad, Residue3D], List[Atom3D]]]:
        if len(self.tetrads) == 0:
            return {}, {}

        ions_channel = defaultdict(list)
        ions_outside = defaultdict(list)

        for ion in self.ions:
            min_distance = math.inf
            min_tetrad = self.tetrads[0]

            for tetrad in self.tetrads:
                distance = numpy.linalg.norm(ion.coordinates() - tetrad.center())
                if distance < min_distance:
                    min_distance = distance
                    min_tetrad = tetrad

            # TODO: verify threshold of 6A between an ion and tetrad channel
            if min_distance < 6.0:
                ions_channel[min_tetrad].append(ion)
                continue

            min_distance = math.inf
            min_tetrad = self.tetrads[0]
            min_nt = min_tetrad.nt1

            for tetrad in self.tetrads:
                for nt in tetrad.nucleotides:
                    for atom in nt.atoms:
                        distance = numpy.linalg.norm(ion.coordinates() - atom.coordinates())
                        if distance < min_distance:
                            min_distance = distance
                            min_tetrad = tetrad
                            min_nt = nt

            # TODO: verify threshold of 3A between an ion and an atom
            if min_distance < 3.0:
                ions_outside[(min_tetrad, min_nt)].append(ion)
                continue

            logging.debug(f'Skipping an ion, because it is too far from any tetrad (distance={min_distance})')

        for tetrad, ions in ions_channel.items():
            tetrad.ions_channel = ions
        for pair, ions in ions_outside.items():
            tetrad, residue = pair
            tetrad.ions_outside[residue] = ions

    def __generate_twoline_dotbracket(self) -> Tuple[str, str, str, Dict[Residue3D, int]]:
        layer1, layer2 = [], []
        for tetrad in self.tetrads:
            layer1.extend([tetrad.pair_12, tetrad.pair_34])
            layer2.extend([tetrad.pair_23, tetrad.pair_41])
        sequence, line1, shifts = self.__elimination_conflicts(layer1)
        _, line2, _ = self.__elimination_conflicts(layer2)
        return sequence, line1, line2, shifts

    def __elimination_conflicts(self, pairs: List[BasePair3D]) -> Tuple[str, str, Dict[Residue3D, int]]:
        orders = dict()
        order = 0
        queue = list(pairs)
        removed = []

        while queue:
            conflicts = defaultdict(list)
            for pi, pj in itertools.combinations(queue, 2):
                if pi.conflicts_with(pj):
                    conflicts[pi].append(pj)
                    conflicts[pj].append(pi)
            if conflicts:
                pair, _ = max(conflicts.items(), key=lambda x: (len(x[1]), x[0].nt1))
                removed.append(pair)
                queue.remove(pair)
            else:
                orders.update({pair: order for pair in queue})
                queue, removed = removed, []
                order += 1

        opening = '([{<' + string.ascii_uppercase
        closing = ')]}>' + string.ascii_lowercase
        dotbracket = dict()
        for pair, order in orders.items():
            nt1, nt2 = sorted([pair.nt1, pair.nt2])
            dotbracket[nt1] = opening[order]
            dotbracket[nt2] = closing[order]

        sequence = ''
        structure = ''
        shifts = dict()
        shift_value = 0
        chain = None
        for nt in sorted(filter(lambda nt: nt.is_nucleotide, self.structure3d.residues), key=lambda nt: nt.index):
            if chain and chain != nt.chain:
                sequence += '-'
                structure += '-'
                shift_value += 1
            sequence += nt.one_letter_name
            structure += dotbracket.get(nt, '.')
            shifts[nt] = shift_value
            chain = nt.chain
        return sequence, structure, shifts

    def __str__(self):
        builder = f'Chain order: {" ".join(self.__chain_order())}\n'
        for helix in self.helices:
            builder += str(helix)
        builder += f'{self.sequence}\n{self.line1}\n{self.line2}'
        return builder

    def __chain_order(self) -> List[str]:
        only_nucleic_acids = filter(lambda nt: nt.is_nucleotide, self.structure3d.residues)
        return list({nt.chain: 0 for nt in sorted(only_nucleic_acids, key=lambda nt: nt.index)}.keys())

    def canonical(self) -> List[BasePair3D]:
        return [base_pair for base_pair in self.base_pairs if base_pair.is_canonical()]

@dataclass
class Visualizer:
    analysis: Analysis
    tetrads: List[Tetrad]
    complete2d: bool
    onz_dict: Dict[BasePair3D, ONZ] = field(init=False)

    def __post_init__(self):
        self.onz_dict = {pair: tetrad.onz for tetrad in self.tetrads for pair in
                         [tetrad.pair_12, tetrad.pair_23, tetrad.pair_34, tetrad.pair_41]}

    def visualize(self, prefix: str, suffix: str):
        fasta = tempfile.NamedTemporaryFile('w+', suffix='.fasta')
        fasta.write(f'>{prefix}-{suffix}\n')
        fasta.write(self.analysis.sequence)
        fasta.flush()

        layer1, layer2 = [], []
        for tetrad in self.tetrads:
            layer1.extend([tetrad.pair_12, tetrad.pair_34])
            layer2.extend([tetrad.pair_23, tetrad.pair_41])
        helix1 = self.__to_helix(layer1, self.analysis.canonical() if self.complete2d else [])
        helix2 = self.__to_helix(layer2)

        currdir = os.path.dirname(os.path.realpath(__file__))
        output_pdf = f'{prefix}-{suffix}.pdf'
        run = subprocess.run([os.path.join(currdir, 'quadraw.R'), fasta.name, helix1.name, helix2.name, output_pdf],
                             stdout=subprocess.PIPE,
                             stderr=subprocess.PIPE)
        if run.returncode == 0:
            print('\nPlot:', output_pdf)
        else:
            logging.error(f'Failed to prepare visualization, reason:\n  {run.stderr.decode()}')

    def __to_helix(self, layer: List[BasePair3D],
                   canonical: Optional[List[BasePair3D]] = None) -> tempfile.NamedTemporaryFile():
        onz_value = {ONZ.O_PLUS: 1, ONZ.O_MINUS: 2, ONZ.N_PLUS: 3, ONZ.N_MINUS: 4, ONZ.Z_PLUS: 5, ONZ.Z_MINUS: 6}
        nucleotides = self.analysis.structure3d.residues
        shifts = self.analysis.shifts

        helix = tempfile.NamedTemporaryFile('w+', suffix='.helix')
        helix.write(f'#{len(self.analysis.sequence) + 1}\n')
        helix.write('i\tj\tlength\tvalue\n')

        for pair in layer:
            x, y = pair.nt1, pair.nt2
            x, y = nucleotides.index(x) + 1 + shifts[x], nucleotides.index(y) + 1 + shifts[y]
            onz = self.onz_dict[pair]
            helix.write(f'{x}\t{y}\t1\t{onz_value.get(onz, 7)}\n')
        if canonical:
            for pair in canonical:
                x, y = pair.nt1, pair.nt2
                x, y = nucleotides.index(x) + 1 + shifts[x], nucleotides.index(y) + 1 + shifts[y]
                helix.write(f'{x}\t{y}\t1\t8\n')

        helix.flush()
        return helix

class AnalysisSimple:
    def __init__(self, structure2d: Structure2D, structure3d: Structure3D):
        self.pairs: List[BasePair3D] = structure3d.base_pairs(structure2d)
        self.graph: Dict[Residue3D, List[Residue3D]] = structure3d.base_pair_graph(structure2d)
        self.pair_dict: Dict[Tuple[Residue3D, Residue3D], BasePair3D] = structure3d.base_pair_dict(structure2d)

    def has_tetrads(self):
        tetrads = set()
        for i in self.graph:
            for j in filter(lambda x: x != i, self.graph[i]):
                for k in filter(lambda x: x not in (i, j), self.graph[j]):
                    for l in filter(lambda x: x not in (i, j, k) and x in self.graph[i], self.graph[k]):
                        if Tetrad.is_valid(i, j, k, l, self.pair_dict):
                            tetrads.add(frozenset([i, j, k, l]))
                        if len(tetrads) > 1:
                            return True
        return False

def center_of_mass(atoms):
    coords = [atom.coordinates() for atom in atoms]
    xs = (coord[0] for coord in coords)
    ys = (coord[1] for coord in coords)
    zs = (coord[2] for coord in coords)
    return numpy.array((sum(xs) / len(coords), sum(ys) / len(coords), sum(zs) / len(coords)))

def eltetrado(structure2d: Structure2D, structure3d: Structure3D, strict: bool, no_reorder: bool,
              stacking_mismatch: int) -> Analysis:
    return Analysis(structure2d, structure3d, strict, no_reorder, stacking_mismatch)

def has_tetrad(structure2d: Structure2D, structure3d: Structure3D) -> bool:
    structure = AnalysisSimple(structure2d, structure3d)
    return structure.has_tetrads()


"""
-------------------------------------------------
   File Name:    Losses.py
   Author:       Zhonghao Huang
   Date:         2019/10/21
   Description:  Module implementing various loss functions
                 Copy from: https://github.com/akanimax/pro_gan_pytorch
-------------------------------------------------
"""

import torch
import torch.nn as nn
from torch.nn import BCEWithLogitsLoss
import numpy as np
from utils import gradient_penalty

# =============================================================
# Interface for the losses
# =============================================================

class GANLoss:
    """ Base class for all losses
        @args:
        disc: Discriminator used for calculating the loss
             Note this must be a part of the GAN framework
    """

    def __init__(self, disc):
        self.disc = disc

    def disc_loss(self, real_samps, fake_samps, alpha, height):
        """
        calculate the discriminator loss using the following data
        :param real_samps: batch of real samples
        :param fake_samps: batch of generated (fake) samples
        :param height: current height at which training is going on
        :param alpha: current value of the fader alpha
        :return: loss => calculated loss Tensor
        """
        raise NotImplementedError("disc_loss method has not been implemented")

    def gen_loss(self, real_samps, fake_samps, alpha, height):
        """
        calculate the generator loss
        :param real_samps: batch of real samples
        :param fake_samps: batch of generated (fake) samples
        :param height: current height at which training is going on
        :param alpha: current value of the fader alpha
        :return: loss => calculated loss Tensor
        """
        raise NotImplementedError("gen_loss method has not been implemented")

class ConditionalGANLoss:
    """ Base class for all conditional losses """

    def __init__(self, disc):
        self.disc = disc

    def disc_loss(self, real_samps, fake_samps, labels, alpha, height):
        raise NotImplementedError("disc_loss method has not been implemented")

    def gen_loss(self, real_samps, fake_samps, labels, alpha, height):
        raise NotImplementedError("gen_loss method has not been implemented")

# =============================================================
# Normal versions of the Losses:
# =============================================================

class StandardGAN(GANLoss):

    def __init__(self, disc):
        super().__init__(disc)
        # define the criterion and activation used for object
        self.criterion = BCEWithLogitsLoss()

    def disc_loss(self, real_samps, fake_samps, alpha, height):
        # small assertion:
        assert real_samps.device == fake_samps.device, \
            "Real and Fake samples are not on the same device"

        # device for computations:
        device = fake_samps.device

        # predictions for real images and fake images separately :
        r_preds = self.disc(real_samps, alpha, height)
        f_preds = self.disc(fake_samps, alpha, height)

        # calculate the real loss:
        real_loss = self.criterion(
            torch.squeeze(r_preds),
            torch.ones(real_samps.shape[0]).to(device))

        # calculate the fake loss:
        fake_loss = self.criterion(
            torch.squeeze(f_preds),
            torch.zeros(fake_samps.shape[0]).to(device))

        # return final losses
        return (real_loss + fake_loss) / 2

    def gen_loss(self, _, fake_samps, alpha, height):
        preds, _, _ = self.disc(fake_samps, alpha, height)
        return self.criterion(
            torch.squeeze(preds),
            torch.ones(fake_samps.shape[0]).to(fake_samps.device))

class WGAN_GP(GANLoss):

    def __init__(self, disc, LAMBDA_GP=10):
        super().__init__(disc)
        self.LAMBDA_GP = LAMBDA_GP

    def disc_loss(self, real_samps, fake_samps, alpha, height):
        critic_real = self.disc(real_samps, alpha, height)
        critic_fake = self.disc(fake_samps, alpha, height)
        gp = gradient_penalty(self.disc, real_samps, fake_samps, alpha, height, device=real_samps.device)
        loss = (
            -(torch.mean(critic_real) - torch.mean(critic_fake))
            + self.LAMBDA_GP * gp
            + (0.001 * torch.mean(critic_real ** 2))
        )
        return loss

    def gen_loss(self, _, fake_samps, alpha, height):
        return -torch.mean(self.disc(fake_samps, alpha, height))

class HingeGAN(GANLoss):

    def __init__(self, disc):
        super().__init__(disc)

    def disc_loss(self, real_samps, fake_samps, alpha, height):
        r_preds = self.disc(real_samps, alpha, height)
        f_preds = self.disc(fake_samps, alpha, height)

        loss = (torch.mean(nn.ReLU()(1 - r_preds)) +
                torch.mean(nn.ReLU()(1 + f_preds)))

        return loss

    def gen_loss(self, _, fake_samps, alpha, height):
        return -torch.mean(self.disc(fake_samps, alpha, height))

class RelativisticAverageHingeGAN(GANLoss):

    def __init__(self, disc):
        super().__init__(disc)

    def disc_loss(self, real_samps, fake_samps, alpha, height):
        # Obtain predictions
        r_preds = self.disc(real_samps, alpha, height)
        f_preds = self.disc(fake_samps, alpha, height)

        # difference between real and fake:
        r_f_diff = r_preds - torch.mean(f_preds)

        # difference between fake and real samples
        f_r_diff = f_preds - torch.mean(r_preds)

        # return the loss
        loss = (torch.mean(nn.ReLU()(1 - r_f_diff))
                + torch.mean(nn.ReLU()(1 + f_r_diff)))

        return loss

    def gen_loss(self, real_samps, fake_samps, alpha, height):
        # Obtain predictions
        r_preds = self.disc(real_samps, alpha, height)
        f_preds = self.disc(fake_samps, alpha, height)

        # difference between real and fake:
        r_f_diff = r_preds - torch.mean(f_preds)

        # difference between fake and real samples
        f_r_diff = f_preds - torch.mean(r_preds)

        # return the loss
        return (torch.mean(nn.ReLU()(1 + r_f_diff))
                + torch.mean(nn.ReLU()(1 - f_r_diff)))

class LogisticGAN(GANLoss):
    def __init__(self, disc):
        super().__init__(disc)

    # gradient penalty
    def R1Penalty(self, real_img, alpha, height):

        # TODO: use_loss_scaling, for fp16
        apply_loss_scaling = lambda x: x * torch.exp(x * torch.Tensor([np.float32(np.log(2.0))]).to(real_img.device))
        undo_loss_scaling = lambda x: x * torch.exp(-x * torch.Tensor([np.float32(np.log(2.0))]).to(real_img.device))

        real_img = torch.autograd.Variable(real_img, requires_grad=True)
        real_logit = self.disc(real_img, alpha, height)
        # real_logit = apply_loss_scaling(torch.sum(real_logit))
        real_grads = torch.autograd.grad(outputs=real_logit, inputs=real_img,
                                         grad_outputs=torch.ones(real_logit.size()).to(real_img.device),
                                         create_graph=True, retain_graph=True)[0].view(real_img.size(0), -1)
        # real_grads = undo_loss_scaling(real_grads)
        r1_penalty = torch.sum(torch.mul(real_grads, real_grads))
        return r1_penalty

    def disc_loss(self, real_samps, fake_samps, alpha, height, r1_gamma=10.0):
        # Obtain predictions
        r_preds = self.disc(real_samps, alpha, height)
        f_preds = self.disc(fake_samps, alpha, height)

        loss = torch.mean(nn.Softplus()(f_preds)) + torch.mean(nn.Softplus()(-r_preds))

        if r1_gamma != 0.0:
            r1_penalty = self.R1Penalty(real_samps.detach(), alpha, height) * (r1_gamma * 0.5)
            loss += r1_penalty

        return loss

    def gen_loss(self, _, fake_samps, alpha, height):
        f_preds = self.disc(fake_samps, alpha, height)

        return torch.mean(nn.Softplus()(-f_preds))