"""
=================
Roman Instruments
=================

WARNING: This model has not yet been validated against other PSF
         simulations, and uses several approximations (e.g. for
         mirror polishing errors, which are taken from HST).
"""

import os.path
import poppy
import numpy as np

from scipy.interpolate import griddata, RegularGridInterpolator
from astropy.io import fits
import astropy.units as u
import logging

from . import utils
from . import webbpsf_core
from .optics import _fix_zgrid_NaNs

_log = logging.getLogger('webbpsf')
import pprint

GRISM_FILTERS = ('GRISM0', 'GRISM1')
PRISM_FILTERS = ('PRISM',)

class WavelengthDependenceInterpolator(object):
    """WavelengthDependenceInterpolator can be configured with
    `n_zernikes` worth of Zernike coefficients at up to `n_wavelengths`
    wavelengths, and will let you `get_aberration_terms` for any
    wavelength in range interpolated linearly between measured/known
    points
    """

    def __init__(self, n_wavelengths=16, n_zernikes=22):
        self._n_wavelengths = n_wavelengths
        self._n_zernikes = n_zernikes
        self._aberration_terms = np.zeros((n_wavelengths, n_zernikes), dtype=np.float64)
        self._wavelengths = []

    def set_aberration_terms(self, wavelength, zernike_array):
        """Supply a reference `wavelength` and a `zernike_array`
        (of length `n_zernikes`) where the aberration is known
        """
        n_wavelengths_set = len(self._wavelengths)
        if wavelength not in self._wavelengths and n_wavelengths_set < self._n_wavelengths:
            self._wavelengths.append(wavelength)
            aberration_row_idx = n_wavelengths_set  # which is now index of last row
        elif wavelength in self._wavelengths:
            aberration_row_idx = self._wavelengths.index(wavelength)
        else:
            # can't add more wavelengths without allocating new _aberration_terms array
            raise ValueError("Already have information at {} wavelengths "
                             "(pass larger n_wavelengths to __init__?)".format(self._n_wavelengths))
        if len(zernike_array) != self._n_zernikes:
            raise ValueError("Expected {} aberration terms (pass different "
                             "n_zernikes to __init__?)".format(self._n_zernikes))
        self._aberration_terms[aberration_row_idx] = zernike_array

    def get_aberration_terms(self, wavelength):
        """Return the Zernike coefficients as interpolated for this
        `wavelength`"""
        # return array of length n_zernikes interpolated for this wavelength
        if wavelength in self._wavelengths:
            # aberration known exactly for this wavelength
            aberration_row_idx = self._wavelengths.index(wavelength)
            return self._aberration_terms[aberration_row_idx]
        else:
            # we have to interpolate @ this wavelength
            aberration_terms = griddata(self._wavelengths, self._aberration_terms, wavelength, method='linear')
            if np.any(np.isnan(aberration_terms)):
                if isinstance(wavelength, u.Quantity):
                    wavelength = wavelength.to(u.m).value
                wavelength_closest = np.clip(wavelength, np.min(self._wavelengths), np.max(self._wavelengths))
                _log.warn("Attempted to get aberrations at wavelength {:.2g} "
                          "outside the range of the reference data; clipping to closest wavelength {:.2g}".format(
                    wavelength, wavelength_closest))

                aberration_terms = griddata(self._wavelengths, self._aberration_terms, wavelength_closest,
                                            method='linear')
            return aberration_terms

class FieldDependentAberration(poppy.ZernikeWFE):
    """FieldDependentAberration incorporates aberrations that
    are interpolated in wavelength, x, and y pixel positions by
    computing the Zernike coefficients for a particular wavelength
    and position.
    """

    """By default, `get_aberration_terms` will zero out Z1, Z2, and Z3
    (piston, tip, and tilt) as they are not meaningful for telescope
    PSF calculations (the former is irrelevant, the latter two would
    be handled by a distortion solution). Change
    `_omit_piston_tip_tilt` to False to include the Z1-3 terms."""
    _omit_piston_tip_tilt = True
    _field_position = None

    def __init__(self, pixel_width, pixel_height,
                 name="Field-dependent Aberration", radius=1.0, oversample=1, interp_order=3):
        self.pixel_width, self.pixel_height = pixel_width, pixel_height
        self.field_position = pixel_width // 2, pixel_height // 2
        self._wavelength_interpolators = {}
        self.pupil_diam = radius * 2.0
        super().__init__(
            name=name,
            verbose=True,
            radius=radius,
            oversample=oversample,
            interp_order=interp_order
        )

    def get_opd(self, wave):
        """Set the Zernike coefficients (for ZernikeWFE.getOPD) based
        on the wavelength of the incoming wavefront and the pixel
        position
        """
        if not isinstance(wave, poppy.Wavefront):
            wavelength = wave
        else:
            wavelength = wave.wavelength
        self.coefficients = wavelength * self.get_aberration_terms(wavelength)
        return super().get_opd(wave)

    @property
    def field_position(self):
        return self._field_position

    @field_position.setter
    def field_position(self, position):
        """Set the x and y pixel position on the detector for which to
        interpolate aberrations"""
        x_pixel, y_pixel = position
        if x_pixel > self.pixel_width or x_pixel < 0:
            raise ValueError("Requested pixel_x position lies outside "
                             "the detector width ({})".format(x_pixel))
        if y_pixel > self.pixel_height or y_pixel < 0:
            raise ValueError("Requested pixel_y position lies outside "
                             "the detector height ({})".format(y_pixel))

        self._field_position = x_pixel, y_pixel

    def add_field_point(self, x_pixel, y_pixel, interpolator):
        """Supply a wavelength-space interpolator for a pixel position
        on the detector"""
        self._wavelength_interpolators[(x_pixel, y_pixel)] = interpolator

    def get_aberration_terms(self, wavelength):
        """Supply the Zernike coefficients for the aberration based on
        the wavelength and pixel position on the detector"""
        if self.field_position in self._wavelength_interpolators:
            # short path: this is a known point
            interpolator = self._wavelength_interpolators[self.field_position]
            coefficients = interpolator.get_aberration_terms(wavelength)
        else:
            # get aberrations at all field points
            field_points, aberration_terms = [], []
            for field_point_coords, point_interpolator in self._wavelength_interpolators.items():
                field_points.append(field_point_coords)
                aberration_terms.append(point_interpolator.get_aberration_terms(wavelength))
            aberration_array = np.asarray(aberration_terms)
            assert len(aberration_array.shape) == 2, "computed aberration array is not 2D " \
                                                     "(inconsistent number of Zernike terms " \
                                                     "at each point?)"
            field_position = tuple(self.field_position)
            coefficients = griddata(
                np.asarray(field_points),
                np.asarray(aberration_terms),
                field_position,
                method='linear'
            )
            if np.any(np.isnan(coefficients)):
                # FIND TWO CLOSEST INPUT GRID POINTS:
                dist = []
                corners = field_points[1:]  # use only the corner points
                for i, ip in enumerate(corners):
                    dist.append(np.sqrt(((ip[0] - field_position[0]) ** 2) + ((ip[1] - field_position[1]) ** 2)))
                min_dist_indx = np.argsort(dist)[:2]  # keep two closest points
                # DEFINE LINE B/W TWO POINTS, FIND ORTHOGONAL LINE AT POINT OF INTEREST,
                # AND FIND INTERSECTION OF THESE TWO LINES.
                x1, y1 = corners[min_dist_indx[0]]
                x2, y2 = corners[min_dist_indx[1]]
                dx = x2 - x1
                dy = y2 - y1
                a = (dy * (field_position[1] - y1) + dx * (field_position[0] - x1)) / (dx * dx + dy * dy)
                closest_interp_point = (x1 + a * dx, y1 + a * dy)
                # INTERPOLATE ABERRATIONS TO CLOSEST INTERPOLATED POINT:
                coefficients = griddata(
                    np.asarray(field_points),
                    np.asarray(aberration_terms),
                    closest_interp_point,
                    method='linear')
                # IF CLOSEST INTERPOLATED POINT IS STILL OUTSIDE THE INPUT GRID,
                # THEN USE NEAREST GRID POINT INSTEAD:
                if np.any(np.isnan(coefficients)):
                    coefficients = aberration_terms[min_dist_indx[0] + 1]
                    _log.warn("Attempted to get aberrations at field point {} which is outside the range "
                              "of the reference data; approximating to nearest input grid point".format(field_position))
                else:
                    _log.warn("Attempted to get aberrations at field point {} which is outside the range "
                              "of the reference data; approximating to nearest interpolated point {}".format(
                        field_position, closest_interp_point))
                assert not np.any(np.isnan(coefficients)), "Could not compute aberration " \
                                                           "at field point {}".format(field_position)
        if self._omit_piston_tip_tilt:
            _log.debug("Omitting piston/tip/tilt")
            coefficients[:3] = 0.0  # omit piston, tip, and tilt Zernikes
        return coefficients

def _load_wfi_detector_aberrations(filename):
    from astropy.io import ascii
    zernike_table = ascii.read(filename, encoding='utf-8-sig')
    detectors = {}

    def build_detector_from_table(number, zernike_table):
        """Build a FieldDependentAberration optic for a detector using
        Zernikes Z1-Z22 at various wavelengths and field points"""
        single_detector_info = zernike_table[zernike_table['sca'] == number]
        field_points = set(single_detector_info['field_point'])
        interpolators = {}
        detector = FieldDependentAberration(
            4096,
            4096,
            radius=RomanInstrument.PUPIL_RADIUS,
            name="Field Dependent Aberration (SCA{:02})".format(number)
        )
        for field_id in field_points:
            field_point_rows = single_detector_info[single_detector_info['field_point'] == field_id]
            local_x, local_y = field_point_rows[0]['local_x'], field_point_rows[0]['local_y']
            interpolator = build_wavelength_dependence(field_point_rows)

            midpoint_pixel = 4096 / 2
            # (local_x in mm / 10 um pixel size) -> * 1e2
            # local_x and _y range from -20.44 to +20.44, so adding to the midpoint pixel
            # makes sense to place (-20.44, -20.44) at (4, 4)
            pixx, pixy = (round(midpoint_pixel - local_x * 1e2),
                          round(midpoint_pixel + local_y * 1e2))

            detector.add_field_point(pixx, pixy, interpolator)
        return detector

    def build_wavelength_dependence(rows):
        """Build an interpolator object that interpolates Z1-Z22 in
        wavelength space"""
        wavelengths = set(rows['wavelength'])
        interpolator = WavelengthDependenceInterpolator(n_wavelengths=len(wavelengths),
                                                        n_zernikes=22)
        for row in rows:
            z = np.zeros(22)
            for idx in range(22):
                z[idx] = row['Z{}'.format(idx + 1)]
            interpolator.set_aberration_terms(row['wavelength'] * 1e-6, z)

        return interpolator

    detector_ids = set(zernike_table['sca'])
    for detid in detector_ids:
        detectors["SCA{:02}".format(detid)] = build_detector_from_table(detid, zernike_table)

    return detectors

@utils.combine_docstrings
class RomanInstrument(webbpsf_core.SpaceTelescopeInstrument):
    PUPIL_RADIUS = 2.4 / 2.0
    """
    RomanInstrument contains data and functionality common to Roman
    instruments, such as setting the pupil shape
    """
    telescope = "Roman"

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.options['jitter'] = 'gaussian'
        self.options['jitter_sigma'] = 0.012 # arcsec/axis, see https://roman.ipac.caltech.edu/sims/Param_db.html#telescope

    def calc_psf(self, outfile=None, source=None, nlambda=None, monochromatic=None,
                 fov_arcsec=None, fov_pixels=None, oversample=None, detector_oversample=None, fft_oversample=None,
                 overwrite=True, display=False, save_intermediates=False, return_intermediates=False,
                 normalize='first', add_distortion=False, crop_psf=False):
        """
        Compute a PSF

        Parameters
        ----------
        add_distortion : bool
            Included for API compatibility with the JWST instrument classes, but has no
            effect on the results for Roman WFI PSF calculations.
        crop_psf : bool
            Included for API compatibility with the JWST instrument classes, but has no
            effect on the results for Roman WFI PSF calculations.

        """

        if add_distortion is not False or crop_psf is not False:
            raise AttributeError('`add_distortion` and `crop_psf` still under '
                                 'development for Roman WFI')
            # return default values to True once implemented

        # Run poppy calc_psf
        psf = webbpsf_core.SpaceTelescopeInstrument.calc_psf(self, outfile=outfile, source=source, nlambda=nlambda,
                                                monochromatic=monochromatic, fov_arcsec=fov_arcsec,
                                                fov_pixels=fov_pixels, oversample=oversample,
                                                detector_oversample=detector_oversample, fft_oversample=fft_oversample,
                                                overwrite=overwrite, display=display,
                                                save_intermediates=save_intermediates,
                                                return_intermediates=return_intermediates, normalize=normalize)

        return psf

    # slightly different versions of the following two functions
    # from the parent superclass
    # in order to interface with the FieldDependentAberration class
    @property
    def detector_position(self):
        """The pixel position in (X, Y) on the detector"""
        return self._detectors[self._detector].field_position

    @detector_position.setter
    def detector_position(self, position):
        # exact copy of superclass function except we save the
        # into a different location.
        try:
            x, y = map(int, position)
        except ValueError:
            raise ValueError("Detector pixel coordinates must be pairs of nonnegative numbers, "
                             "not {}".format(position))
        if x < 0 or y < 0:
            raise ValueError("Detector pixel coordinates must be nonnegative integers")
        if x > self._detector_npixels - 1 or y > self._detector_npixels - 1:
            raise ValueError("The maximum allowed detector pixel "
                             "coordinate value is {}".format(self._detector_npixels - 1))

        self._detectors[self._detector].field_position = (int(position[0]), int(position[1]))

    def _get_aberrations(self):
        """Get the OpticalElement that applies the field-dependent
        optical aberrations. (Called in get_optical_system.)"""
        return self._detectors[self._detector]

    def _get_fits_header(self, result, options):
        """Populate FITS Header keywords"""
        super()._get_fits_header(result, options)
        result[0].header['DETXPIXL'] = (self.detector_position[0],
                                        'X pixel position (for field dependent aberrations)')
        result[0].header['DETYPIXL'] = (self.detector_position[1],
                                        'Y pixel position (for field dependent aberrations)')
        result[0].header['DETECTOR'] = (self.detector, 'Detector selected')

class WFIPupilController:
    """
    This is a helper class for the WFI and is used to swap in
    the correct pupil each time the detector is changed.
    The pupil depends on pupil_mask, detector, and filter;
    pupil_mask is set automatically upon the receipt of the
    detector and filter values selected by the user in the WFI class.
    The user should not interact with this class directly, only
    through the API provided through the WFI class.
    """

    def __init__(self):
        self._datapath = None
        self._pupil_basepath = None

        self._pupil = None
        self._pupil_mask = None # new for webbpsf 1.0

        # Flag to en-/disable automatic selection of the appropriate pupil_mask
        self._auto_pupil = True

        # Flag to en-/disable automatic selection of the appropriate pupil file
        self._auto_pupil_mask = True

        # Save formattable pupil file names for each pupil mask
        self.pupil_file_formatters = {
            'SKINNY': 'RST_WIM_Filter_skinny_{0}.fits.gz',
            'WIDE': 'RST_WIM_Filter_F184_{0}.fits.gz',
            'GRISM': 'RST_WSM_Grism_Grism_{0}.fits.gz',
            'PRISM': 'RST_WSM_Prism_Prism_{0}.fits.gz'}

    @property
    def pupil(self):
        """
        The path to the FITS file containing pupil information for the
        detector/filter combination sent from the WFI class. Cannot be
        directly set by the user.
        """
        return self._pupil

    @pupil.setter
    def pupil(self, value):
        raise AttributeError('Pupil cannot be directly specified. '
                             'Use lock_pupil() instead.')

    @property
    def pupil_mask(self):
        """
        The corresponding mask for the filter sent from the WFI class.
        (See WFI.pupil_mask_list for a list of valid filters.) Cannot
        be directly set by the user.
        """
        return self._pupil_mask

    @pupil_mask.setter
    def pupil_mask(self, name):
        raise AttributeError('Pupil mask cannot be directly specified. '
                             'Use lock_pupil_mask() instead.')

    def _get_filter_mask(self, wfi_filter):
        """
        Returns the appropriate mask for a given WFI filter.

        Parameters
        ----------
        wfi_filter : string
            See WFI.filter_list for a list of valid filters.
        """
        wfi_filter = wfi_filter.upper()

        if wfi_filter in GRISM_FILTERS:
            return 'GRISM'
        elif wfi_filter in PRISM_FILTERS:
            return 'PRISM'
        elif wfi_filter in ['F184', 'F213']:
            return 'WIDE'
        else:
            # this method should only be called after WFI.filter was validated,
            # so we assume all inputs are valid and direct those that don't pass
            # preceding cases to skinny
            return 'SKINNY'

    def set_base_path(self, datapath):
        """
        Sets the root directory of the path to WebbPSF's data files.
        This should be set before this class is used.

        Parameters
        ----------
        datapath : string
            Path to WebbPSF-WFI data files
        """
        self._datapath = datapath
        self._pupil_basepath = os.path.join(self._datapath, "pupils")

    def update_pupil(self, filter, detector):
        """
        Selects the specific pupil file corresponding with a detector
        and filter combination sent from the WFI class. Also finds and
        indirectly sets the proper pupil_mask in the process.

        Parameters
        ----------
        filter : string
            See WFI.filter_list for a list of valid filters.

        detector : string
            See WFI.detector_list for a list of valid detectors.
        """
        if not self._auto_pupil:
            _log.info('Automatic pupil selection was locked; '
                      'using user-provided pupil.')
            return

        if self._pupil_basepath is None:
           raise Exception('update_pupil called before setting pupil file path')

        # change detector string to match file format (e.g., "SCA01" -> "SCA_1")
        det_substr = f"{detector[:3]}_{str(int((detector[3:])))}"

        # figure out proper mask based on filter (or use locked mask if enabled)
        pupil_mask = (self._get_filter_mask(filter) if self._auto_pupil_mask
                      else self.pupil_mask)

        path_formatter = self.pupil_file_formatters[pupil_mask]
        pupil = os.path.join(self._pupil_basepath,
                             path_formatter.format(det_substr))

        self._pupil_mask = pupil_mask
        self._pupil = pupil

        _log.info(f"Using {'' if self._auto_pupil_mask else 'locked '}"
                  f"pupil mask '{pupil_mask}' and detector '{detector}'.")

    def lock_pupil(self, pupil_path):
        """
        Prevents the WFIPupilController class from dynamically updating
        the path to the pupil on any changes to the detector or filter
        selected in the WFI class. Instead, the path remains locked on
        whichever `pupil_path` was provided to this method.

        CAUTION: This is non-standard usage of the WFI class and may
        lead to unexpected behavior.

        Parameters
        ----------
        pupil_path : string
            The custom path to your pupil file.
        """
        self._pupil_mask = None
        self._pupil = pupil_path
        self._auto_pupil = False

    def unlock_pupil(self):
        """
        Undoes the effects of lock_pupil() and resets WFIPupilController
        to its default state of updating the pupil whenever a detector
        or filter is changed in the WFI class.
        """
        self._auto_pupil = True

    def lock_pupil_mask(self, pupil_mask):
        """
        Prevents the WFIPupilController class from dynamically updating
        the pupil mask on any changes to the filter selected in the WFI
        class. Instead, the pupil mask remains locked on whichever
        `pupil_mask` was provided to this method.

        CAUTION: This is non-standard usage of the WFI class and may
        lead to unexpected behavior.

        Parameters
        ----------
        filter : string
            See WFI.pupil_mask_list for a list of valid pupil masks.
        """
        if pupil_mask not in self.pupil_file_formatters.keys():
            raise Exception('invalid pupil mask')
        elif not self._auto_pupil:
            raise Exception('Pupil is locked. Unlock pupil before locking pupil mask.')
        else:
            self._pupil_mask = pupil_mask
            self._auto_pupil_mask = False

    def unlock_pupil_mask(self):
        """
        Undoes the effects of lock_pupil_mask() and resets
        WFIPupilController to its default state of updating the pupil
        mask whenever filter is changed in the WFI class.
        """
        self._auto_pupil_mask = True

class WFI(RomanInstrument):
    """
    WFI represents the Roman mission's Wide Field Imager.

    WARNING: This model has not yet been validated against other PSF
             simulations, and uses several approximations (e.g. for
             mirror polishing errors, which are taken from HST).
    """

    def __init__(self):
        # pixel scale is from Roman-AFTA SDT report final version (p. 91)
        # https://roman.ipac.caltech.edu/sims/Param_db.html
        pixelscale = 110e-3 # arcsec/px

        # Initialize the aberrations for super().__init__
        self._aberration_files = {}
        self._is_custom_aberration = False
        self._current_aberration_file = ""

        super().__init__("WFI", pixelscale=pixelscale)

        # Initialize the pupil controller
        self._pupil_controller = WFIPupilController()
        self._pupil_controller.set_base_path(self._datapath)

        self.pupil_mask_list = list(self._pupil_controller.pupil_file_formatters.keys())

        # Define default aberration files for WFI modes
        self._aberration_files = {
            'imaging': os.path.join(self._datapath, 'wim_zernikes_cycle9.csv'),
            'prism': os.path.join(self._datapath,
                                  'wsm_prism_zernikes_cycle9.csv'),
            'grism': os.path.join(self._datapath,
                                  'wsm_grism_zernikes_cycle9.csv'),
            'custom': None}

        # Load and set default detector from aberration file
        self._detector_npixels = 4096
        self._load_detector_aberrations(self._aberration_files[self.mode])
        self.detector = 'SCA01'

        self.opd_list = [os.path.join(self._WebbPSF_basepath,
                                      'upscaled_HST_OPD.fits')]
        self.pupilopd = self.opd_list[-1]

    def _addAdditionalOptics(self, optsys, **kwargs):
        return optsys, False, None

    def _load_detector_aberrations(self, path):
        """
        Helper function that, given a path to a file containing detector
        aberrations, loads the Zernike values and populates the class'
        dictator list with `FieldDependentAberration` detectors. This
        function achieves this by calling the
        `webbpsf.roman._load_wfi_detector_aberrations` function.

        Users should use the `override_aberrations` function to override
        current aberrations.

        Parameters
        ----------
        path : string
            The path to the file containing detector aberrations.
        """
        detectors = _load_wfi_detector_aberrations(path)
        assert len(detectors.keys()) > 0

        self._detectors = detectors
        self._current_aberration_file = path

    def _validate_config(self, **kwargs):
        """
        Validates that the WFI is configured sensibly.

        This mainly consists of selecting the masked or unmasked pupil
        appropriately based on the wavelengths requested.
        """
        assert self.filter is not None, 'filter is None'
        assert self.detector is not None, 'detector is None'
        self._update_pupil()

        assert self.pupil is not None, 'pupil is None'
        super()._validate_config(**kwargs)

    def _get_filter_mode(self, wfi_filter):
        """
        Given a filter name, returns the WFI mode.

        Parameters
        ----------
        wfi_filter : string
            Name of WFI filter. See WFI.filter_list for valid values.

        Returns
        -------
        mode : string
            Returns 'imaging', 'grism', or 'prism' depending on filter.

        Raises
        ------
        ValueError
            ...if the input filter is not found in the WFI filter list.
        """

        wfi_filter = wfi_filter.upper()
        if wfi_filter in GRISM_FILTERS:
            return 'grism'
        elif wfi_filter in PRISM_FILTERS:
            return 'prism'
        elif wfi_filter in self.filter_list:
            return 'imaging'
        else:
            raise ValueError(f"Instrument {self.name} doesn't have a filter "
                             f"called {wfi_filter}.")

    def _update_pupil(self, filter=None, detector=None):
        if detector is None:
            detector = self.detector
        if filter is None:
            filter = self.filter

        if detector is not None and filter is not None:
            self._pupil_controller.update_pupil(filter=filter,detector=detector)

    @RomanInstrument.detector.setter
    def detector(self, value):
        """
        The current WFI detector. See WFI.detector_list for valid values.
        """
        if value.upper() not in self.detector_list:
            raise ValueError("Invalid detector. Valid detector names are: {}".format(', '.join(self.detector_list)))

        self._detector = value.upper()
        if self._detector is not None:
            self._update_pupil(detector=self._detector)

    @RomanInstrument.filter.setter
    def filter(self, value):
        """
        The current WFI filter. See WFI.filter_list for valid values.
        """
        # Update filter
        value = value.upper()

        if value not in self.filter_list:
            raise ValueError(f"Instrument {self.name} doesn't have a "
                             f"filter called {value}.")

        self._filter = value

        # Update aberrations if self._aberration_files has been initiated (not
        # empty) and if they haven't been locked by user
        if self._aberration_files and not self._is_custom_aberration:

            # identify aberration file for new mode
            mode = self._get_filter_mode(self._filter)
            aberration_file = self._aberration_files[mode]

            # if aberrations are not already loaded for the new mode,
            # load and replace detectors using the new mode's aberration file
            if not os.path.samefile(self._current_aberration_file,
                                    aberration_file):
                self._load_detector_aberrations(aberration_file)

        # Update pupil only if detector was previously loaded
        # ( i.e., skip this step when called by super() )
        if self.detector is not None:
            self._update_pupil(filter=self._filter)

    @property
    def pupil(self):
        """
        The path to the FITS file containing pupil information for the
        detector/filter combination sent from the WFI class. Cannot be
        directly set by the user.
        """
        return self._pupil_controller.pupil

    @pupil.setter
    def pupil(self, value):
        # don't allow pupil to be set until the pupil controller is active. (a
        # parent class tries to set it to None in WFI's preceding super() call)
        if hasattr(self, '_pupil_controller'):
            raise AttributeError('Pupil cannot be directly specified. '
                                 'Use lock_pupil() instead.')

    @property
    def pupil_mask(self):
        """
        The corresponding mask for the current filter. Cannot be
        directly set by the user.
        """
        return self._pupil_controller.pupil_mask

    @pupil_mask.setter
    def pupil_mask(self, name):
        raise AttributeError('Pupil mask cannot be directly specified. '
                             'Use lock_pupil_mask() instead.')

    @property
    def mode(self):
        """
        The current WFI mode. Cannot be directly set by the user.
        """
        return self._get_filter_mode(self.filter)

    @mode.setter
    def mode(self, value):
        raise AttributeError("WFI mode cannot be directly specified; "
                             "it is set by changing filters.")

    def lock_aberrations(self, aberration_path):
        """
        This function loads user provided aberrations from a file and
        locks this instrument to only use the provided aberrations (even
        if the filter or mode change).

        To release the lock and load the default aberrations, use
        unlock_aberrations(). To load new user provided
        aberrations, call this function with the new path.

        To load custom aberrations, please provide a csv file
        containing the detector names, field point positions and Zernike
        values. The file should contain the following column
        names/values (comments in parentheses should not be included):
            - sca (Detector number)
            - wavelength (µm)
            - field_point (field point number/ID for SCA and wavelength,
            starts with 1)
            - local_x (mm, local detector coords)
            - local_y (mm, local detector coords)
            - global_x (mm, global instrument coords)
            - global_y (mm, global instrument coords)
            - axis_local_angle_x (XAN)
            - axis_local_angle_y (YAN)
            - wfe_rms_waves (nm)
            - wfe_pv_waves (waves)
            - Z1 (Zernike phase NOLL coefficients)
            - Z2 (Zernike phase NOLL coefficients)
            - Z3 (Zernike phase NOLL coefficients)
            - Z4 (Zernike phase NOLL coefficients)
              .
              .
              .

        Please refer to the default aberration files for examples. If
        you have the WebbPSF data installed and defined, you can get the
        path to that file by running the following:
            >>> from webbpsf import roman
            >>> wfi = roman.WFI()
            >>> print(wfi._aberration_files["imaging"])

        Warning: You should not edit the default files!
        """
        self._load_detector_aberrations(aberration_path)
        self._aberration_files['custom'] = aberration_path
        self._is_custom_aberration = True

    def unlock_aberrations(self):
         """
         Releases the lock on the detector aberration file location
         and loads the default file.
         """
         aberration_path = self._aberration_files[self.mode]
         self._load_detector_aberrations(aberration_path)
         self._aberration_files['custom'] = None
         self._is_custom_aberration = False

    def lock_pupil(self, pupil_path):
        """
        Prevents dynamic updates of the path to the proper pupil file on
        any changes to the selected detector or filter. Instead, the
        path remains locked on whichever `pupil_path` was provided here.

        WARNING: This is non-standard usage of the WFI class and may
        lead to unexpected behavior.

        Parameters
        ----------
        pupil_path : string
            The custom path to your pupil file.
        """
        if os.path.isfile(pupil_path):
            self._pupil_controller.lock_pupil(pupil_path)
        else:
            raise FileNotFoundError(f"{pupil_path} not found.")

        _log.warning("Disabling default pupil selection behavior.")

    def unlock_pupil(self):
        """
        Undoes the effects of lock_pupil() by resetting the class to
        its default state of updating the pupil whenever a detector or
        filter is changed. If necessary, it also sets the proper pupil
        for the current detector/filter combination.
        """
        self._pupil_controller.unlock_pupil()
        self._update_pupil() # reset pupil

    def lock_pupil_mask(self, pupil_mask):
        """
        Prevents dynamic updates of the pupil mask on any change to the
        selected filter. Instead, the pupil mask remains locked on
        whichever `pupil_mask` was provided here.

        WARNING: This is non-standard usage of the WFI class and may
        lead to unexpected behavior.

        Parameters
        ----------
        filter : string
            See WFI.pupil_mask_list for a list of valid pupil masks.
        """
        self._pupil_controller.lock_pupil_mask(pupil_mask)
        self._update_pupil()

    def unlock_pupil_mask(self):
        """
        Undoes the effects of lock_pupil_mask() and resets the class to
        its default state of updating the pupil mask whenever the filter
        is changed.
        """
        self._pupil_controller.unlock_pupil_mask()
        self._update_pupil() # reset pupil mask

class RomanCoronagraph(RomanInstrument):
    """
    Roman Coronagraph Instrument

    Simulates the PSF of the Roman coronagraph.

    Current functionality is limited to the Shaped Pupil Coronagraph (SPC)
    observing modes, and these modes are only simulated with static, unaberrated
    wavefronts, without relay optics and without DM control. The design
    respresented here is an approximation to a baseline concept, and will be
    subject to change based on trades studies and technology development.

    Parameters
    ----------
    mode : str
        Roman Coronagraph Instrument observing mode. If not specified, the
        __init__ function will set this to a default mode 'CHARSPC_F660'
    pixelscale : float
        Detector pixelscale. If not specified, the pixelscale will default to
        0.02 arcsec for configurations usint the IMAGER camera and 0.025 arcsec
        for the IFS.
    fov_arcsec : float
        Field of view in arcseconds. If not specified, the field of view will
        default to 3.20 arcsec for the IMAGER camera and 1.76 arcsec for the IFS.

    """

    camera_list = ['IMAGER', 'IFS']
    filter_list = ['F660', 'F721', 'F770', 'F890']
    apodizer_list = ['CHARSPC', 'DISKSPC']
    fpm_list = ['CHARSPC_F660_BOWTIE', 'CHARSPC_F770_BOWTIE', 'CHARSPC_F890_BOWTIE', 'DISKSPC_F721_ANNULUS']
    lyotstop_list = ['LS30D88']

    _mode_table = {  # MODE             CAMERA    FILTER  APODIZER   FPM             LYOT STOP
        'CHARSPC_F660': ('IFS', 'F660', 'CHARSPC', 'CHARSPC_F660_BOWTIE', 'LS30D88'),
        'CHARSPC_F770': ('IFS', 'F770', 'CHARSPC', 'CHARSPC_F770_BOWTIE', 'LS30D88'),
        'CHARSPC_F890': ('IFS', 'F890', 'CHARSPC', 'CHARSPC_F890_BOWTIE', 'LS30D88'),
        'DISKSPC_F721': ('IMAGER', 'F721', 'DISKSPC', 'DISKSPC_F721_ANNULUS', 'LS30D88')}

    def __init__(self, mode=None, pixelscale=None, fov_arcsec=None, apply_static_opd=False):
        super().__init__("RomanCoronagraph", pixelscale=pixelscale)

        self._detector_npixels = 1024
        self._detectors = {camera: 'placeholder' for camera in self.camera_list}

        self.pupil_mask_list = self.lyotstop_list  # alias for use in webbpsf_core
        self.image_mask_list = self.fpm_list  # alias for use in webbpsf_core
        self.pupil = os.path.join(self._WebbPSF_basepath, 'AFTA_CGI_C5_Pupil_onax_256px_flip.fits')
        if apply_static_opd:
            self.pupilopd = os.path.join(self._WebbPSF_basepath, 'CGI', 'OPD', 'CGI_static_OPD.fits')
        else:
            self.pupilopd = None
        self.aberration_optic = None
        self.options = {'force_coron': True}
        # Allow the user to pre-emptively override the default instrument FoV and pixel scale
        if fov_arcsec is not None:
            self.fov_arcsec = fov_arcsec
            self._override_fov = True
        else:
            self._override_fov = False
        if pixelscale is not None:
            self._pixelscale = pixelscale
            self._override_pixelscale = True
        else:
            self._override_pixelscale = False

        if mode is None:
            self.print_mode_table()
            _log.info("Since the mode was not specified at instantiation, defaulting to CHARSPC_F660")
            self.mode = 'CHARSPC_F660'
        else:
            self.mode = mode

    @property
    def camera(self):
        """Currently selected camera name"""
        return self._camera

    @camera.setter
    def camera(self, value):
        value = value.upper()  # force to uppercase
        if value not in self.camera_list:
            raise ValueError("Instrument {0} doesn't have a camera called {1}.".format(self.name, value))
        self._camera = value
        if value == 'IMAGER':
            if not hasattr(self, 'fov_arcsec') or not self._override_fov:
                self.fov_arcsec = 3.2
            if not hasattr(self, 'pixelscale') or not self._override_pixelscale:
                self.pixelscale = 0.020  # Nyquist at 465 nm
        else:  # default to 'IFS'
            if not hasattr(self, 'fov_arcsec') or not self._override_fov:
                self.fov_arcsec = 2 * 0.82  # 2015 SDT report, Section 3.4.1.1.1:
                                            # IFS has 76 lenslets across the (2 x 0.82) arcsec FoV.
            if not hasattr(self, 'pixelscale') or not self._override_pixelscale:
                self.pixelscale = 0.025  # Nyquist at 600 nm

    # for coronagraph, there is one detector per camera and it should be set automatically.
    @property
    def detector(self):
        return self.camera

    @detector.setter
    def detector(self, value):
        raise RuntimeError("Can't set detector directly for RomanCoronagraph; set camera instead.")

    @property
    def filter(self):
        """Currently selected filter name"""
        return self._filter

    @filter.setter
    def filter(self, value):
        value = value.upper()  # force to uppercase
        if value not in self.filter_list:
            raise ValueError("Instrument {0} doesn't have a filter called {1}.".format(self.name, value))
        self._filter = value

    @property
    def apodizer(self):
        """Currently selected apodizer name"""
        return self._apodizer

    @apodizer.setter
    def apodizer(self, value):
        value = value.upper()  # force to uppercase
        if value not in self.apodizer_list:
            raise ValueError("Instrument {0} doesn't have a apodizer called {1}.".format(self.name, value))
        self._apodizer = value
        if value == 'DISKSPC':
            self._apodizer_fname = \
                os.path.join(self._datapath, "optics/DISKSPC_SP_256pix.fits.gz")
        else:  # for now, default to CHARSPC
            self._apodizer_fname = \
                os.path.join(self._datapath, "optics/CHARSPC_SP_256pix.fits.gz")

    @property
    def fpm(self):
        """Currently selected FPM name"""
        return self._fpm

    @fpm.setter
    def fpm(self, value):
        value = value.upper()  # force to uppercase
        if value not in self.fpm_list:
            raise ValueError("Instrument {0} doesn't have a FPM called {1}.".format(self.name, value))
        self._fpm = value
        if value.startswith('DISKSPC'):
            self._fpmres = 3
            self._owa = 20.
            self._Mfpm = int(np.ceil(self._fpmres * self._owa))
            self._fpm_fname = \
                os.path.join(self._datapath,
                             "optics/DISKSPC_FPM_65WA200_360deg_-_FP1res{0:d}_evensamp_D{1:03d}_{2:s}.fits.gz".format(
                                 self._fpmres, 2 * self._Mfpm, self.filter))
        else:
            self._fpmres = 4
            self._owa = 9.
            self._Mfpm = int(np.ceil(self._fpmres * self._owa))
            self._fpm_fname = \
                os.path.join(self._datapath,
                             "optics/CHARSPC_FPM_25WA90_2x65deg_-_FP1res{0:d}_evensamp_D{1:03d}_{2:s}.fits.gz".format(
                                 self._fpmres, 2 * self._Mfpm, self.filter))

    @property
    def lyotstop(self):
        """Currently selected Lyot stop name"""
        return self._lyotstop

    @lyotstop.setter
    def lyotstop(self, value):
        # preserve case for this one since we're used to that with the lyot mask names
        if value not in self.lyotstop_list:
            raise ValueError("Instrument {0} doesn't have a Lyot mask called {1}.".format(self.name, value))
        self._lyotstop = value
        self._lyotstop_fname = os.path.join(self._datapath, "optics/SPC_LS_30D88_256pix.fits.gz")

    @property
    def mode_list(self):
        """Available Observation Modes"""
        keys = self._mode_table.keys()
        keys = sorted(keys)
        return keys

    # mode works differently since it's a meta-property that affects the other ones:
    @property
    def mode(self):
        """Currently selected mode name"""
        for modename, settings in self._mode_table.items():
            if (self.camera == settings[0].upper() and self.filter == settings[1].upper() and
                    self.apodizer == settings[2].upper() and self.fpm == settings[3].upper() and
                    self.lyotstop == settings[4]):
                return modename
        return 'Custom'

    @mode.setter
    def mode(self, value):
        if value not in self.mode_list:
            raise ValueError("Instrument {0} doesn't have a mode called {1}.".format(self.name, value))
        settings = self._mode_table[value]
        self.camera = settings[0]
        self.filter = settings[1]
        self.apodizer = settings[2]
        self.fpm = settings[3]
        self.lyotstop = settings[4]
        _log.info('Set the following optical configuration:')
        _log.info('camera = {0}, filter = {1}, apodizer = {2}, fpm = {3}, lyotstop = {4}'.format(\
                  self.camera, self.filter, self.apodizer, self.fpm, self.lyotstop))

    def print_mode_table(self):
        """Print the table of observing mode options and their associated optical configuration"""
        _log.info("Printing the table of Roman Coronagraph Instrument observing modes supported by WebbPSF.")
        _log.info("Each is defined by a combo of camera, filter, apodizer, "
                  "focal plane mask (FPM), and Lyot stop settings:")
        _log.info(pprint.pformat(self._mode_table))

    @property
    def detector_position(self):
        """The pixel position in (X, Y) on the detector"""
        return 512, 512

    @detector_position.setter
    def detector_position(self, position):
        raise RuntimeError("Detector position not adjustable for RomanCoronagraph")

    def _validate_config(self, **kwargs):
        super()._validate_config(**kwargs)

    def _addAdditionalOptics(self, optsys, oversample=4):
        """Add coronagraphic or spectrographic optics for RomanCoronagraph."""

        trySAM = False

        if ('pupil_shift_x' in self.options and self.options['pupil_shift_x'] != 0) or \
                ('pupil_shift_y' in self.options and self.options['pupil_shift_y'] != 0):
            shift = (self.options['pupil_shift_x'], self.options['pupil_shift_y'])
        else:
            shift = None

        # Add the shaped pupil apodizer
        optsys.add_pupil(transmission=self._apodizer_fname, name=self.apodizer, shift=None)

        # Add the FPM
        optsys.add_image(transmission=self._fpm_fname, name=self.fpm)

        # Add Lyot stop
        self.pupil_mask = self.lyotstop
        optsys.add_pupil(transmission=self._lyotstop_fname, name=self.lyotstop, shift=shift)

        # Cast as MatrixFTCoronagraph; this configures the detector
        occ_box_size = 1.
        mft_optsys = poppy.MatrixFTCoronagraph(optsys, oversample=oversample, occulter_box=occ_box_size)

        return mft_optsys, trySAM, occ_box_size

    def _get_aberrations(self):
        """Get the OpticalElement that applies the field-dependent
        optical aberrations. (Called in get_optical_system.)"""
        return None

    def _get_fits_header(self, result, options):
        """Populate FITS Header keywords"""
        super()._get_fits_header(result, options)
        pupil_hdr = fits.getheader(self.pupil)
        apodizer_hdr = fits.getheader(self._apodizer_fname)
        fpm_hdr = fits.getheader(self._fpm_fname)
        lyotstop_hdr = fits.getheader(self._lyotstop_fname)

        result[0].header.set('MODE', self.mode, comment='Observing mode')
        result[0].header.set('CAMERA', self.camera, comment='Imager or IFS')
        result[0].header.set('APODIZER', self.apodizer, comment='Apodizer')
        result[0].header.set('APODTRAN', os.path.basename(self._apodizer_fname),
                             comment='Apodizer transmission')
        result[0].header.set('PUPLSCAL', apodizer_hdr['PUPLSCAL'],
                             comment='Apodizer pixel scale in m/pixel')
        result[0].header.set('PUPLDIAM', apodizer_hdr['PUPLDIAM'],
                             comment='Full apodizer array size, incl padding.')
        result[0].header.set('FPM', self.fpm, comment='Focal plane mask')
        result[0].header.set('FPMTRAN', os.path.basename(self._fpm_fname),
                             comment='FPM transmission')
        result[0].header.set('FPMSCAL', fpm_hdr['PIXSCALE'], comment='FPM spatial sampling, arcsec/pix')
        result[0].header.set('LYOTSTOP', self.lyotstop, comment='Lyot stop')
        result[0].header.set('LSTRAN', os.path.basename(self._lyotstop_fname),
                             comment='Lyot stop transmission')
        result[0].header.set('PUPLSCAL', lyotstop_hdr['PUPLSCAL'],
                             comment='Lyot stop pixel scale in m/pixel')
        result[0].header.set('PUPLDIAM', lyotstop_hdr['PUPLDIAM'],
                             comment='Lyot stop array size, incl padding.')

import numpy,from scipy,import astropy,from astropy


import numpy as np
import matplotlib.pyplot as plt
import h5py as h5
from . import tof_tools
import os,sys

__all__ = ["sample",'get_montrig','count_rate','dtc','inc']

class sample:
    """
    Store all relevant information for a sample being
    studied at the RPI Linac.

    Stores and calculates data about a sample studied at
    the RPI Linac and follows the file structure and naming
    conventions that are found in codes originating from
    Brian McDermott, such as "40mprocess" and "h5query".
    The output from these codes is used as input for this
    class.

    Parameters
    ----------
    samp_name : str
        Name of sample of interest, e.g. Ta-2mm, MUST be
        the same as it appears in file structure of initial
        raw binary files as it is typically propagated
        through to the HDF5 file.
    data_type : str
        Type of data which was collected, e.g. yield,
        or transmission
    tof_arr : numpy array
        Time of flight data, 1st column is ascending
        time bins [us], 2nd column is bin width, 3rd
        column is the summed counts
    hdf_address : str, optional
        File address of the hdf file output from 40mprocess
        which contains: monitors, triggers, and pulse events;
        must end with slash
    phys_dims : array_like, optional
        Physical dimensions [cm] of sample being studied, must be
        in format [thickness,length,width] or [thickness,radius]
        where thickness is the length of sample parallel to the
        Linac beam
    tau : float, optional
        The dead time of the system [us] used to collect the counts
        data, default value is 0
    paralyzeable : bool, optional
        If system used to collect counts data was paralyzeable
        a different dead time model must be used, default is
        non-paralyzeable dead-time model
    plot_montrig : bool, optional
        Plot the monitor ratio through the montrig function
    factor : 1-d numpy array, optional
        Compression factors for the tof groups, 1st compression factor
        corresponds to time of flights below c_pt[0], 2nd corresponds to tofs above
        c_pt[0] and below c_pt[1], and so on until end of range is found.
    c_pts : 1-d numpy array, optional
        Compression time of flights that seperate the bins that are to be grouped.
        List points from lowest to highest.
    avg : bool, optional
        Return an average of group values instead of
        the sum in the group. group error is 1/N*(sum(error^2))^(1/2)
    binedge : bool, optional
        Return TOF column for grouped data as the first
        tof in group instead of the mean tof of the
        group.
    verbose : bool, optional
        Choose whether to print out approximate comp pt energy locations
        and indices. Default is False.
    FP : float, optional
        Flight path needed to convert the time of flight values of bins to
        energies for the compression point selection.
    verbose : bool, optional
        Option to print out extra information

    Notes
    -----
    Will eventually require break points and grouping
    metadata.

    Allows for yield calculations, must later include
    transmission.
    """
    def __init__(self,samp_name,data_type,tof_arr,hdf_address=None,
                 phys_dims=None,tau=0.0,paralyzeable=False,
                 plot_montrig=False,factor=None,c_pts=None,
                 avg=False,binedge=False,FP=None,verbose=False):

        if verbose == True:
            print("--------\n {} data for {}\n--------".format(data_type,
                  samp_name))

        self.samp_name = samp_name
        self.phys_dims = phys_dims
        self.data_type = data_type

        #------------------------------------------------------------
        #------------ Monitors and triggers -------------------------
        #------------------------------------------------------------
        default_montrig = False

        if hdf_address == None:

            default_montrig = True
            print("----------\n----------")
            print("Monitors and triggers not given, cps calculated with")
            print("default 1e8 monitor and trigger counts.")
            print("----------\n----------")

            self.montrig = np.array([1e8,1e8])

        else:
            self.montrig = get_montrig(hdf_address,samp_name,plot=plot_montrig,
                                       verbose=verbose)

        #------------------------------------------------------------
        #------------ Count rate ------------------------------------
        #------------------------------------------------------------
        if default_montrig:
            self.cps = count_rate(tof_arr,self.montrig[1],tau,paralyzeable=paralyzeable,
                 factor=factor,c_pts=c_pts,avg=avg,binedge=binedge,verbose=verbose,FP=FP)
        else:
            self.cps = count_rate(tof_arr,self.montrig[1][1],tau,paralyzeable=paralyzeable,
                 factor=factor,c_pts=c_pts,avg=avg,binedge=binedge,verbose=verbose,FP=FP)
        #------------------------------------------------------------

#------------------------------------------------------------------------------

def get_montrig(hdf_address,samp_name,plot=False,verbose=False):
    """
    Collect monitors and triggers used to reduce
    capture data to proper dead time corrected
    monitor normalized count rates

    Parameters
    ----------
    hdf_address : str
        File address of the HDF5 file, must end with "/"
    samp_name : str
        Name of the sample being analyzed. Should
        match the sample name in the HDF5 filename
    plot : bool, optional
        Choose whether to plot the monitors and triggers
        as a function of cycle number
    verbose : bool, optional
        Choose whether to print out trigger and monitor
        information.

    Returns
    -------
    montrig : numpy array
        Vector of monitor counts and triggers

    Notes
    -----
    The montrig array is printed out when function
    is called, user can see order of monitors and
    triggers.
    """

    with h5.File(hdf_address+samp_name+"_master.h5","r") as hdf:

        mon_counts = []
        det_counts = []
        sum_mon = np.zeros(8)
        sum_trig = np.zeros(2)

        for cycle in hdf[samp_name]:
            mon_vector = hdf[samp_name+"/"+cycle+"/MON"]
            trig_vector = hdf[samp_name+"/"+cycle+"/TRIG"][0]
            datag_array = hdf[samp_name+"/"+cycle+"/DATA/GOOD"]
            datat_array = hdf[samp_name+"/"+cycle+"/DATA/TRUNC"]
            sum_detcts = len(datag_array)+len(datat_array)
            mon_counts.append(mon_vector)
            det_counts.append(sum_detcts)

            # number of monitors stored in /MON != to 8
            if (len(sum_mon) != len(mon_vector)) or (len(sum_trig) != len(trig_vector)):
                raise ValueError("Unexpected number of mon's in /MON")

            sum_mon += mon_vector
            sum_trig += trig_vector

        num_cyc = len(hdf[samp_name])
        mon_counts = np.transpose(mon_counts)
        # monitor ratios
        mr0,mr1,mr2 = np.zeros(num_cyc),np.zeros(num_cyc),np.zeros(num_cyc)
        for i in range(num_cyc):
            mr0[i] = det_counts[i]/mon_counts[0][i]
            mr1[i] = det_counts[i]/mon_counts[1][i]
            mr2[i] = det_counts[i]/mon_counts[2][i]
        # std of monitor cycles / mean of mon. cycles
        stdmean = [0,0,0]
        stdmean[0] = np.std(mr0)/np.mean(mr0)
        stdmean[1] = np.std(mr1)/np.mean(mr1)
        stdmean[2] = np.std(mr2)/np.mean(mr2)

        if plot:
            x = range(num_cyc)
            plt.plot(x,mr0/np.mean(mr0),label="Sum/Mon 0")
            plt.plot(x,mr1/np.mean(mr1),label="Sum/Mon 1")
            plt.plot(x,mr2/np.mean(mr2),label="Sum/Mon 2")
            plt.legend()
            plt.show()

        montrig = np.array([sum_mon,sum_trig])
        if verbose == True:
            print("         Counts          Std/mean")
            for i in range(len(stdmean)):
                print("Mon. {} : {}  ,  {}".format(i,sum_mon[i],stdmean[i]))
            if len(sum_trig>1):
                print("Triggers : {}  {}".format(montrig[1][0],montrig[1][1]))

    return montrig

def count_rate(tof_arr,trig,tau=None,paralyzeable=False,factor=None,c_pts=None,
               avg=False,binedge=False,verbose=False,FP=None):
    """
    Converts counts to dead-time corrected count
    rates

    Parameters
    ----------
    tof_arr : 2-d numpy array
        Time of flight data, 1st column is ascending
        time bins [us], 2nd column is the summed counts,
        and the 3rd column (if desired) is error on the
        counts. 3rd column only necessary if avg==True.
        !! time of flight data **MUST** be equal bin sizes !!
    trig : float
        Triggers recorded for the sample
    tau : float
        The dead time of the system used to collect
        counts data [us]
    paralyzeable : bool,optional
        If the system used to collect the counts data
        was paralyzeable, a different dead time model
        must be used
    factor : 1-d numpy array, optional
        Compression factors for the tof groups, 1st compression factor
        corresponds to time of flights below c_pt[0], 2nd corresponds to tofs above
        c_pt[0] and below c_pt[1], and so on until end of range is found.
    c_pts : 1-d numpy array
        Compression time of flights that seperate the bins that are to be grouped.
        List points from lowest to highest.
    avg : bool, optional
        Return an average of group values instead of
        the sum in the group. group error is 1/N*(sum(error^2))^(1/2)
    binedge : bool, optional
        Return TOF column for grouped data as the first
        tof in group instead of the mean tof of the
        group.
    verbose : bool, optional
        Choose whether to print out approximate comp pt energy locations
        and indices. Default is False.
    FP : float
        Flight path needed to convert the time of flight values of bins to
        energies for the compression point selection.

    Returns
    -------
    dc_arr : 2-d numpy array
        TOF data, 1st column ascending time bins [us],
        2nd column is count rate [cps], 3rd column is
        the error on that count rate

    Notes
    -----
    Units matter!!

    Error for count rates are sqrt of the counts vector user
    provides.

    Dead time correction error given in the dead time correction
    function `dtc()`.

    .. math:: CPS = \\frac{C}{b_w Trig}

    .. math:: \\Delta C = \\sqrt{C}

    .. math:: \\Delta CPS = \\sqrt{\\left(\\frac{\\Delta C}{b_w Trig}\\right)^2}

    .. math:: CPS_{DC} = (CPS) dtcf

    .. math:: \\Delta CPS_{DC}=\\sqrt{\\left(\\frac{\\Delta CPS}{CPS} \\right)^2 + \\left(\\frac{\\Delta dtcf}{dtcf} \\right)^2}

    """
    use_single_group = False
    use_comp_group = False

    # -------------------------------------------------------------------------
    # ---- run some checks ----------------------------------------------------
    # -------------------------------------------------------------------------
    # did the user give one grouping input but not the other?
    mismatch1 = ((factor is None) and (c_pts is not None))
    mismatch2 = ((factor is not None) and (c_pts is None))
    if mismatch1:
        raise ValueError("c_pts given, but factor is not.")
    # if mismatch2, we may need to use single group if factor is not an array
    if mismatch2:
        try:
            len(factor)
            raise ValueError()
        except TypeError:
            print('Using single group')
            use_single_group = True
        except:
            print('Factor is an array and c_pts is not given. Aborting.')
            raise ValueError("Check input for argument: factor.")
    if factor is not None and c_pts is not None:
        use_comp_group = True
    if use_comp_group:
        # we need the flight path for comp_group
        if FP is None:
            raise ValueError("Specify the flight path for comp_group()")

    # cast array into np array
    tof_arr = np.array(tof_arr)
    # if input is pandas, may need to transpose
    if tof_arr.shape[0] > tof_arr.shape[1]:
        tof_arr = np.transpose(tof_arr)

    # -------------------------------------------------------------------------
    # ---- Calculate and apply dead time correction ---------------------------
    # -------------------------------------------------------------------------
    if tau is None:
        tau = 0.0     # will be "corrected" by 1.0, with error 0.0

    tau = tau*1e-6    # [us --> s]

    if paralyzeable == False:
        bin_width = tof_arr[0][1]-tof_arr[0][0]
        dtcf = dtc(tof_arr[1],tau,trig,bin_width)
        dc_counts = tof_arr[1]*dtcf[0]
        # sum of the squares of relative errors in counts and dead-time corr.
        # time the corrected counts
        if avg:
            ecounts = tof_arr[2]
        else:
            ecounts = np.sqrt(tof_arr[1])

        rel_err_dtcf = (dtcf[1]/dtcf[0])
        rel_err_counts = (ecounts/tof_arr[1])
        edc_counts = np.sqrt(rel_err_dtcf**2 + rel_err_counts**2) * dc_counts
        # --------------------------------------------
        # check if error on correction is significant
        # --------------------------------------------
        ratio = rel_err_dtcf/(rel_err_dtcf+rel_err_counts)
        if (ratio > 0.02).any():
            print("Max % of dead time error: ",max(ratio))
            print("!! Dead time correction factor accounts for > 2% error in the counts !!")
            print("!! This is not accounted for in grouping if avg is not used !!")
            print("To fix, you have to change the hard-coded grouping function.")

    else:
        raise NotImplementedError("paralyzeable dead time not implemented.")

    # -------------------------------------------------------------------------
    # ---- Group the counts if user has specified -----------------------------
    # -------------------------------------------------------------------------
    if use_single_group:
        # include the input error on the counts
        gtof_array = tof_tools.single_group([tof_arr[0],dc_counts,ecounts],
            factor,avg=avg,binedge=binedge)
    if use_comp_group:
        # include the input error on the counts
        gtof_array = tof_tools.comp_group([tof_arr[0],dc_counts,ecounts],
            factor,c_pts,FP,avg=avg,binedge=binedge,verbose=verbose)

    # -------------------------------------------------------------------------
    # ---- Calculate the count rate -------------------------------------------
    # -------------------------------------------------------------------------
    if use_comp_group is False and use_single_group is False:
        cps = dc_counts/bin_width/trig/1e-6
        # error in the cps
        err_cps = 1/bin_width/trig/1e-6 * edc_counts
        # stage a tof array for the return statement
        time_of_flight = tof_arr[0]
    else:
        cps = gtof_array[2]/gtof_array[1]/trig/1e-6
        # error in cps
        err_cps = 1/gtof_array[1]/trig/1e-6 * gtof_array[3]
        # stage a tof array for the return statement
        time_of_flight = gtof_array[0]

    cps_arr = np.array([time_of_flight,cps,err_cps])

    return cps_arr

def dtc(counts,dead_time,trigs,bin_width):
    """
    Calculate the TOF dead-time correction factor according
    to the non-paralyzable model

    Parameters
    ----------
    counts : array_like
        A 1-d array of counts as a function of tof (tof
        in ascending order, equal bin spacing throughout)
    dead_time : float
        The dead time [us] associated with the data collection
        system.
    trigs : int
        The number of triggers = the number of times channel
        was collecting data. For RPI the number of LINAC
        pulses.
    bin_width : float
        The width of the bins in time [us]. Should be constant
        for the array of counts.

    Returns
    -------
    dtcf : 2-d numpy array
        The dead time correction factor for the counts array
        based on the non-paralyzeable dead time model is in
        column 0 of array dtcf, the error associated with it
        is in column 1.
        len(dtcf[0]) = len(dtcf[1]) = len(counts).

    Notes
    -----
    dtc() returns 0's for the first bins that span the dead
    time. 1's would keep the signal closer to reality, but
    0's make it obvious to the user that those channels do
    not have a correction applied, and data in this time of
    flight range is rarely used since it usually occurs before
    the gamma flash.

    .. math:: dtcf_{\\tau>b_w} = \\frac{1}{1-\\frac{SUM}{trig}} = \\frac{1}{1-\\frac{\\sum_iw_iC_i}{trig}}

    where the weights are all 1 except the first and last bin.
    First weight is 0.5, last weight is the fraction of bins
    the dead time minus 0.5 covers (e.g. 3.1 minus 1/2 = 2.6)
    minus the whole bins dead time minus 0.5 covers (2), so
    the last weight is 0.6.

    Now the error.

    .. math:: \\Delta SUM = \\sqrt{\\sum_i(w_i\\sqrt{C_i})^2}

    .. math:: \\Delta dtcf_{\\tau>b_w} = \\sqrt{\\left(\\frac{\\partial dtcf_{\\tau>b_w}}{\\partial SUM}\\Delta SUM\\right)^2} = \\sqrt{\\left(\\frac{\\Delta SUM}{trig\\left(1-\\frac{SUM}{trig}\\right)^2}\\right)^2}

    .. math:: = \\frac{\\Delta SUM}{trig\\left(1-\\frac{SUM}{trig}\\right)^2} = \\frac{\\Delta SUM}{trig}dtcf_{\\tau>b_w}^2

    If the dead time is less than the bin width:

    .. math:: dtcf_{\\tau<b_w} = \\frac{1}{1-\\frac{C\\tau}{b_wtrig}}

    .. math:: \\Delta dtcf_{\\tau<b_w} = \\sqrt{\\left(\\frac{\\partial dtcf_{\\tau<b_w}}{\\partial C}\\Delta C\\right)^2} = \\frac{\\tau\\sqrt{C}}{b_wtrig\\left(1-\\frac{C\\tau}{b_wtrig}\\right)^2}

    Algorithm is a Pythonized version of the dead-time-correction algorithm by Y. Danon, "Design and Construction of the RPI Enhanced Thermal Neutron Target and Thermal Cross Section Measurements of Rare Earth Isotopes.", Doctoral Thesis, (1993).
    """

    # cast as numpy = major speedup
    counts = np.array(counts)
    lc = len(counts)
    dtcf = np.zeros(lc)
    ddtcf = np.zeros(lc)
    if bin_width>=dead_time:
        x = counts*dead_time/bin_width/trigs
        dtcf = 1/(1-x)
        ddtcf = dead_time*np.sqrt(counts)/bin_width/trigs/(1-x)**2
    else:
        # how many bins does the dead time span - 0.5 a bin
        an = dead_time/bin_width-0.5
        # how many whole bins does the dead time span
        n = int(an)
        # bins 0-n will be left 0
        print("Channels 0 -",n,"set to 0")
        # fraction of earliest bin that dead time spans
        f1 = an-n
        # add half of current bin
        SUM = counts[n+1:lc] * 0.5
        # (w_i(C_i)^0.5)^2 = C_i*0.5*0.5
        sSUM = SUM*0.5
        # add preceding whole bins within span of dead time
        for j in range(1,n+1):
            SUM += counts[n+1-j:lc-j]
            sSUM += counts[n+1-j:lc-j]
        # add the fraction of the last bin dead time spans
        SUM += f1*counts[0:lc-n-1]
        sSUM += f1**2*counts[0:lc-n-1]
        dSUM = np.sqrt(sSUM)

        dtcf[n+1:lc] = 1/(1-SUM/trigs)
        ddtcf[n+1:lc] = dSUM/trigs*dtcf[n+1:lc]**2
    return np.array([dtcf,ddtcf])

def inc(x):
    return x + 1


#    This file is part of qdpy.
#
#    qdpy is free software: you can redistribute it and/or modify
#    it under the terms of the GNU Lesser General Public License as
#    published by the Free Software Foundation, either version 3 of
#    the License, or (at your option) any later version.
#
#    qdpy is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
#    GNU Lesser General Public License for more details.
#
#    You should have received a copy of the GNU Lesser General Public
#    License along with qdpy. If not, see <http://www.gnu.org/licenses/>.

"""This file contains the CMA-ME emitters, as defined in paper "Covariance Matrix Adaptation for the Rapid Illumination of Behavior Space"(https://arxiv.org/abs/1912.02400). Additional algorithms are also included, notably the ME-MAP-Elites algorithm (Multi-Emitter MAP-Elites), as presented in the paper "Multi-Emitter MAP-Elites: Improving quality, diversity and convergence speed with heterogeneous sets of emitters".  All emitters are implemented as subclasses of the CMA-ES class, so the library "cma" needs to be installed to use them. """

__all__ = ["CMAMEOptimizingEmitter", "CMAMERandomDirectionEmitter", "CMAMEImprovementEmitter", "MEMAPElitesUCB1"]

from typing import Optional, Tuple, List, Iterable, Iterator, Any, TypeVar, Generic, Union, Sequence, MutableSet, MutableSequence, Type, Callable, Generator, Mapping, MutableMapping, overload
import numpy as np
import random

from .base import *
from .evolution import *
from qdpy.utils import *
from qdpy.phenotype import *
from qdpy.base import *
from qdpy.containers import *
from qdpy import tools

try:
    import cma
except ImportError:
    pass

@registry.register
class CMAMEOptimizingEmitter(CMAES):
    """TODO"""

#    def __init__(self, container: Container, budget: int,
#            dimension: int, **kwargs):
#        super().__init__(container, budget, dimension=dimension, **kwargs)

    def reinit(self) -> None:
        super().reinit()
        self.restart()

    def restart(self):
        if len(self.container):
            init_ind = random.choice(self.container)
        else:
            init_ind = [0.] * self.dimension
        self.es = cma.CMAEvolutionStrategy(init_ind, self.sigma0, self._opts)
        self._pop_inds.clear()
        self._pop_fitness_vals.clear()

    def _stop(self) -> bool:
        term_cond = self.es.stop(check=False)
        #if len(term_cond) > 0:
        #    print(f"RESTART!!!: {term_cond}")
        return len(term_cond) > 0

    def _internal_ask(self, base_ind: IndividualLike) -> IndividualLike:
        if self._stop():
            self.restart()
        return super()._internal_ask(base_ind)

@registry.register
class CMAMERandomDirectionEmitter(CMAMEOptimizingEmitter):
    """TODO"""

    def restart(self):
        super().restart()
        features_domain = np.array(self.container.features_domain)
        self._direction = np.random.normal(0., 1., features_domain.shape[0])
        self._direction = features_domain[:,0] + self._direction * (features_domain[:,1] - features_domain[:,0])
        self._features_mean = np.zeros(features_domain.shape[0])

    def _internal_tell(self, individual: IndividualLike, added_to_container: bool, xattr: Mapping[str, Any] = {}) -> None:
        if self.ignore_if_not_added_to_container and not added_to_container:
            return

        self._features_mean += np.array(individual.features)
        if added_to_container:
            self._pop_inds += [individual]

        if len(self._pop_inds) >= self.es.popsize:
            self._features_mean /= self.es.popsize
            _pop_delta = []
            for ind in self._pop_inds:
                dv = ind.features - self._features_mean
                _pop_delta.append(self._direction.dot(dv))
            try:
                self.es.tell(self._pop_inds, _pop_delta)
            except RuntimeError:
                pass
            else:
                self._pop_inds.clear()
                self._pop_fitness_vals.clear()

@registry.register
class CMAMEImprovementEmitter(CMAMEOptimizingEmitter):
    """TODO"""

    _novel_pop_inds: MutableSequence[IndividualLike]
    _novel_pop_fitness_vals: MutableSequence[Any]

    def restart(self):
        super().restart()
        self._novel_pop_inds = []
        self._novel_pop_fitness_vals = []

    def _tell_before_container_update(self, individual: IndividualLike) -> Tuple[MutableMapping[str, Any], bool]:
        novel = None
        delta = 0.
        prev = None
        if hasattr(self.container, "index_grid") and hasattr(self.container, "solutions"):
            try:
                idx = self.container.index_grid(individual.features) # type: ignore
                novel = len(self.container.solutions[idx]) == 0 # type: ignore
                if not novel:
                    prev = self.container.solutions[0] # type: ignore
            except:
                pass
        delta = sum(individual.fitness.values)
        if prev is not None:
            delta -= sum(prev.fitness.values)
        return {"novel": novel, "delta": delta}, True

    def _internal_tell(self, individual: IndividualLike, added_to_container: bool, xattr: Mapping[str, Any] = {}) -> None:
        if self.ignore_if_not_added_to_container and not added_to_container:
            return
        novel = xattr.get("novel", False)
        delta = xattr.get("delta", 0.)

        if added_to_container:
            if novel:
                self._novel_pop_inds += [individual]
                self._novel_pop_fitness_vals += [-1. * delta]
            else:
                self._pop_inds += [individual]
                self._pop_fitness_vals += [-1. * delta]

        if len(self._pop_inds) + len(self._novel_pop_inds) >= self.es.popsize:
            tot_pop = []
            tot_pop_fitness_vals = []

            # Sort individuals
            sorted_pop_inds = [self._pop_inds[k] for k in argsort(self._pop_fitness_vals)]
            #sorted_pop_fitness_vals = sorted(self._pop_fitness_vals)
            sorted_novel_pop_inds = [self._novel_pop_inds[k] for k in argsort(self._novel_pop_fitness_vals)]
            #sorted_novel_pop_fitness_vals = sorted(self._novel_pop_fitness_vals)
            for i, ind in enumerate(sorted_novel_pop_inds):
                tot_pop.append(ind)
                tot_pop_fitness_vals.append(i)
            for i, ind in enumerate(sorted_pop_inds):
                tot_pop.append(ind)
                tot_pop_fitness_vals.append(len(sorted_novel_pop_inds) + i)

            try:
                self.es.tell(tot_pop, tot_pop_fitness_vals)
            except RuntimeError:
                pass
            else:
                self._pop_inds.clear()
                self._pop_fitness_vals.clear()
                self._novel_pop_inds.clear()
                self._novel_pop_fitness_vals.clear()

@registry.register
class MEMAPElitesUCB1(AlgWrapper):
    """Implementation of the ME-MAP-Elites algorithm (Multi-Emitter MAP-Elites), as presented in the paper "Multi-Emitter MAP-Elites: Improving quality, diversity and convergence speed with heterogeneous sets of emitters". UCB1 is used to select the active emitters. """
    zeta: float
    nb_active_emitters: int
    _active_emitters_idx: Sequence[int]
    current_active_idx: int
    initial_expected_rwds: float
    shuffle_emitters: bool

    def __init__(self, algorithms: Any, zeta: float = 0.0005, nb_active_emitters: int = 1, initial_expected_rwds: float = 1.,
            shuffle_emitters: bool = True, **kwargs):
        self.zeta = zeta
        self.nb_active_emitters = nb_active_emitters
        self.initial_expected_rwds = initial_expected_rwds
        self.shuffle_emitters = shuffle_emitters
        super().__init__(algorithms, **kwargs)

    def reinit(self) -> None:
        super().reinit()
        self._up_active_emitters()
        self.current_active_idx = 0

    def _up_active_emitters(self) -> None:
        """Update the list of active emitters."""
        #expected_rwds = [
        #        self.initial_expected_rwds if alg.nb_evaluations == 0 else
        #        float(alg.nb_updated) / float(alg.nb_evaluations) + self.zeta * math.sqrt(math.log(self.nb_evaluations) / float(alg.nb_evaluations))
        #        for alg in self.algorithms]
        #self._active_emitters_idx = argsort(expected_rwds, reverse=True)[:self.nb_active_emitters]

        algos_idx = list(range(len(self.algorithms)))
        if self.shuffle_emitters:
            # Shuffle algorithm list
            random.shuffle(algos_idx)
        # Compute expected rewards
        expected_rwds = [
                self.initial_expected_rwds if self.algorithms[a].nb_evaluations == 0 else
                    float(self.algorithms[a].nb_updated) / float(self.algorithms[a].nb_evaluations) +
                    self.zeta * math.sqrt(math.log(self.nb_evaluations) / float(self.algorithms[a].nb_evaluations))
                for a in algos_idx]
        # Update active emitters indexes list
        self._active_emitters_idx = [algos_idx[a] for a in argsort(expected_rwds, reverse=True)][:self.nb_active_emitters]
        #print(f"DEBUG _up_active_emitters: {argsort(expected_rwds, reverse=True)}")
        #print(f"DEBUG _up_active_emitters: {expected_rwds}")
        #print(f"DEBUG _up_active_emitters: {self._active_emitters_idx}")

    def next(self) -> None:
        """Switch to the next algorithm in Sequence `self.algorithms`, if there is one."""
        if self.current_active_idx < len(self._active_emitters_idx) - 1:
            self.current_active_idx += 1
        else:
            self._up_active_emitters()
            self.current_active_idx = 0
        next_idx = self._active_emitters_idx[self.current_active_idx]
        self.switch_to(next_idx)
        #print(f"NEXT: {next_idx}")

# MODELINE	"{{{1
# vim:expandtab:softtabstop=4:shiftwidth=4:fileencoding=utf-8
# vim:foldmethod=marker


import os
import numpy as np
import tensorflow as tf

from .models.schnet import SchNet

def get_atom_indices(n_atoms, batch_size):
    n_distances = n_atoms ** 2 - n_atoms
    seg_m = np.repeat(range(batch_size), n_atoms).astype(np.int32)
    seg_i = np.repeat(np.arange(n_atoms * batch_size), n_atoms - 1).astype(np.int32)
    idx_ik = seg_i
    idx_j = []
    for b in range(batch_size):
        for i in range(n_atoms):
            for j in range(n_atoms):
                if j != i:
                    idx_j.append(j + b * n_atoms)

    idx_j = np.hstack(idx_j).ravel().astype(np.int32)
    offset = np.zeros((n_distances * batch_size, 3), dtype=np.float32)
    ratio_j = np.ones((n_distances * batch_size,), dtype=np.float32)
    seg_j = np.arange(n_distances * batch_size, dtype=np.int32)

    seg_m, idx_ik, seg_i, idx_j, seg_j, offset, ratio_j = \
        tf.constant(seg_m), tf.constant(idx_ik), tf.constant(seg_i), tf.constant(idx_j), \
        tf.constant(seg_j), tf.constant(offset), tf.constant(ratio_j)
    idx_jk = idx_j
    return seg_m, idx_ik, seg_i, idx_j, idx_jk, seg_j, offset, ratio_j

class SchNetMD:
    def __init__(self, energy_model_path, force_model_path=None, batch_size=1,
                 nuclear_charges=6. * np.ones((20,), dtype=np.int64)):

        self.n_atoms = len(nuclear_charges)
        seg_m, idx_ik, seg_i, idx_j, idx_jk, seg_j, offset, ratio_j = get_atom_indices(self.n_atoms, batch_size)

        self.energy_model = self.load_model(energy_model_path)
        self.force_model = self.load_model(force_model_path) \
            if force_model_path is not None else None

        self.positions = tf.placeholder(tf.float32, shape=(batch_size * self.n_atoms, 3))
        self.charges = tf.tile(tf.constant(nuclear_charges.ravel(), dtype=tf.int64),
                               (batch_size,))

        g = tf.get_default_graph()
        with g.gradient_override_map({"Tile": "TileDense"}):
            self.energy = self.energy_model(self.charges, self.positions,
                                            offset, idx_ik, idx_jk, idx_j,
                                            seg_m, seg_i, seg_j, ratio_j)
            if self.force_model is None:
                self.forces = -tf.reshape(tf.convert_to_tensor(
                    tf.gradients(tf.reduce_sum(self.energy), self.positions)[0]),
                    (batch_size, self.n_atoms, 3))
            else:
                energy = self.force_model(self.charges, self.positions,
                                          offset, idx_ik, idx_jk, idx_j,
                                          seg_m, seg_i, seg_j, ratio_j)
                self.forces = -tf.reshape(tf.convert_to_tensor(tf.gradients(tf.reduce_sum(energy),
                                                                            self.positions)[0]),
                                          (batch_size, self.n_atoms, 3))
        self.error = tf.reduce_max(tf.sqrt(tf.reduce_sum(self.forces ** 2, 2)))

        ckpt = tf.train.latest_checkpoint(os.path.join(energy_model_path, 'validation'))
        self.session = tf.Session()
        self.energy_model.restore(self.session, ckpt)
        if self.force_model is not None:
            ckpt = tf.train.latest_checkpoint(os.path.join(force_model_path, 'validation'))
            self.force_model.restore(self.session, ckpt)

    def load_model(self, model_path):
        args = np.load(os.path.join(model_path, 'args.npy')).item()

        model = SchNet(args.interactions, args.basis, args.filters, args.cutoff,
                       intensive=args.intensive, filter_pool_mode=args.filter_pool_mode)
        return model

    def get_energy_and_forces(self, positions):
        positions = positions.reshape((-1, 3)).astype(np.float32)
        feed_dict = {
            self.positions: positions
        }
        E, F = self.session.run([self.energy, self.forces], feed_dict=feed_dict)
        return E, F

    def relax(self, positions, eps=0.01, rate=1e-4):
        err = 100.
        positions = positions.reshape((-1, 3)).astype(np.float32)
        print('Start relaxation')
        count = 0
        while err > eps:
            feed_dict = {
                self.positions: positions
            }
            F, err = self.session.run([self.forces, self.error], feed_dict=feed_dict)
            positions += rate * F[0]
            count += 1
            if count % 100 == 0:
                print('Iteration ', str(count), 'Max Force:', err)
        print('Maximal force length: ' + str(err))
        return positions


#!/usr/bin/env python

# Built-in
import os
import argparse

# Common
import numpy as np

_save = True
_here = os.path.abspath(os.path.dirname(__file__))
_Exp, _Cls, _name = os.path.split(__file__)[1].split('_')[:3]
assert not any([any([ss in s for ss in ['Notes','.']])
               for s in [_Exp, _Cls, _name]])

def get_notes():
    notes = {'DPhi':{}, 'dPhi':{}}
    # Total toiroidal width in equatorial plane
    notes['DPhi'] = 346.756
    # Toroidal gap between central tiles in equatorial plane
    notes['dPhi'] = 1.500
    # Total vertical height
    notes['DZ'] = 1047.517
    # Poloidal gap between central tiles
    notes['dl'] = 2.249
    # Radial width
    notes['DR'] = 134.678

    # sampleXZY
    notes['sampleXZY'] = [[2159.632, 523.758, 2030.517],#0
                          [2149.423, 514.993, 2020.128],
                          [2148.227, 508.130, 2018.912],
                          [2154.559, 442.235, 2071.259],    # h
                          [2166.380, 414.998, 2085.893],    # h
                          [2178.412, 385.686, 2095.532],    # h
                          [2206.053, 371.041, 2077.756],    # e1
                          [2210.863, 369.335, 2082.651],
                          [2221.072, 378.101, 2093.040],
                          [2222.233, 374.874, 2094.221],
                          [2211.288, 368.155, 2083.083],#10
                          [2209.454, 361.586, 2081.216],    # e2
                          [2192.106, 347.606, 2109.467],    # h
                          [2201.248, 317.292, 2118.769],    # h
                          [2210.389, 286.979, 2128.072],    # h
                          [2234.025, 280.107, 2106.220],    # e1
                          [2230.026, 271.564, 2115.303],    # i
                          [2234.631, 268.948, 2119.988],
                          [2245.923, 275.667, 2130.785],
                          [2246.768, 272.245, 2131.645],
                          [2234.940, 267.696, 2120.303],#20
                          [2232.501, 261.541, 2117.821],    # i
                          [2240.819, 252.580, 2113.134],    # e2
                          [2220.353, 246.610, 2138.212],    # h
                          [2226.542, 214.813, 2144.510],    # h
                          [2232.731, 183.016, 2150.808],    # h
                          [2257.455, 167.114, 2130.063],    # e1
                          [2261.769, 163.634, 2134.453],
                          [2273.250, 168.183, 2146.136],
                          [2273.763, 164.630, 2146.658],
                          [2261.957, 162.335, 2134.644],#30
                          [2258.958, 156.705, 2131.592],    # e2
                          [2238.784, 141.097, 2156.967],    # h
                          [2241.907, 108.399, 2160.145],    # h
                          [2245.030, 075.701, 2163.323],    # h
                          [2268.232, 059.603, 2141.030],    # e1
                          [2272.178, 055.324, 2145.045],
                          [2283.984, 057.619, 2157.059],
                          [2284.156, 054.000, 2157.234],
                          [2272.241, 054.000, 2145.109],
                          [2268.736, 049.000, 2141.543],#40 e2
                          [2247.062, 033.000, 2165.391],    # h
                          [2272.659, 011.500, 2137.688],    # i
                          [2247.062, 000.000, 2165.391],    # h
                          [2247.062,-033.000, 2165.391],    # h
                          [2268.736,-049.000, 2141.543],    # e1
                          [2272.241,-054.000, 2145.109],
                          [2284.156,-054.000, 2157.234],
                          [2283.984,-057.619, 2157.059],
                          [2272.178,-055.324, 2145.045],
                          [2268.232,-059.603, 2141.030],#50 e2
                          [2245.030,-075.701, 2163.323],    # h
                          [2241.907,-108.399, 2160.145],    # h
                          [2238.784,-141.097, 2156.967],    # h
                          [2258.958,-156.705, 2131.592],    # e1
                          [2261.957,-162.335, 2134.644],
                          [2273.763,-164.630, 2146.658],
                          [2273.250,-168.183, 2146.136],
                          [2261.769,-163.634, 2134.453],
                          [2257.455,-167.114, 2130.063],#59 e2
                          [2232.731,-183.016, 2150.808],    # h
                          [2226.542,-214.813, 2144.510],    # h
                          [2220.353,-246.610, 2138.212],    # h
                          [2240.819,-252.580, 2113.134],    # e1
                          [2232.501,-261.541, 2117.821],    # i
                          [2234.940,-267.696, 2120.303],
                          [2246.768,-272.245, 2131.645],
                          [2245.923,-275.667, 2130.785],
                          [2234.631,-268.948, 2119.988],
                          [2230.026,-271.564, 2115.303],    # i
                          [2234.025,-280.107, 2106.220],#70 e2
                          [2210.389,-286.979, 2128.072],    # h
                          [2201.248,-317.292, 2118.769],    # h
                          [2192.106,-347.606, 2109.467],    # h
                          [2209.454,-361.586, 2081.216],    # e1
                          [2211.288,-368.155, 2083.083],
                          [2222.233,-374.874, 2094.221],
                          [2221.072,-378.101, 2093.040],
                          [2210.863,-369.335, 2082.651],
                          [2206.053,-371.041, 2077.756],#79 e2
                          [2178.412,-385.686, 2095.532],    # h
                          [2166.486,-413.961, 2083.395],    # h
                          [2154.559,-442.235, 2071.259],    # h
                          [2148.227,-508.130, 2018.912],
                          [2149.423,-514.993, 2020.128],
                          [2159.632,-523.758, 2030.517],
                          [2152.686,-542.612, 2243.735],    # Back
                          [2452.439,-057.045, 2300.172],
                          [3550.000,-057.150, 0000.000],
                          [3550.000, 057.150, 0000.000],
                          [2452.118, 056.737, 2302.204],#90
                          [2151.698, 544.667, 2242.729]]
    notes['sampleXZY'] = np.array(notes['sampleXZY'])
    notes['ind_h'] = [3,4,5,12,13,14,23,24,25,32,33,34,41,43,44,51,52,53,
                      60,61,62,71,72,73,80,81,82]
    notes['ind_Back'] = [86,87,88,89,90,91]
    notes['ind_i'] = [16,21,42,64,69]
    notes['ind_e1'] = [6,15,26,35,45,54,63,74]
    notes['ind_e2'] = [11,22,31,40,50,59,70,79]

    for kk in notes.keys():
        if type(notes[kk]) is dict:
            notes[kk]['In'] = notes[kk]['In']*1.e-3
            notes[kk]['Out'] = notes[kk]['Out']*1.e-3
        elif not 'nb' in kk and not 'ind' in kk:
            notes[kk] = notes[kk]*1.e-3
    return notes

def _get_intersect(D0,u0,D1,u1):
    k = -np.cross(D0-D1,u1)/np.cross(u0,u1)
    return D0 + k*u0

def make_Poly(save=_save, path=_here):
    notes = get_notes()

    Poly = np.array([np.hypot(notes['sampleXZY'][:,0],notes['sampleXZY'][:,2]),
                     notes['sampleXZY'][:,1]])
    # Finish V2
    ind = np.zeros((Poly.shape[1],),dtype=bool)
    ind[notes['ind_h']] = True
    ind[notes['ind_i']] = True
    nind = np.arange(0,Poly.shape[1])
    Polybis = Poly.copy()
    for ii in range(0,len(notes['ind_e1'])+1):
        i0 = notes['ind_e2'][ii-1] if ii>0 else 2
        i1 = notes['ind_e1'][ii] if ii<len(notes['ind_e1']) else 83
        indi = ind & (nind>i0) & (nind<i1)
        D0 = Poly[:,i0]
        u = Poly[:,i1]-D0
        un2 = np.linalg.norm(u)**2
        k = -np.sum((D0[:,np.newaxis]-Poly[:,indi])*u[:,np.newaxis],axis=0)/un2
        Polybis[:,indi] = D0[:,np.newaxis] + k[np.newaxis,:]*u[:,np.newaxis]

    # Make V0 and V1
    ind0 = np.ones((Poly.shape[1],),dtype=bool)
    ind0[notes['ind_h']] = False
    ind0[notes['ind_i']] = False
    ind0[[0,1,84,85]] = False
    Poly0 = Poly[:,ind0]
    inde1 = np.zeros((Poly.shape[1],),dtype=bool)
    inde2 = np.zeros((Poly.shape[1],),dtype=bool)
    inde1[notes['ind_e1']] = True
    inde2[notes['ind_e2']] = True
    inde1 = inde1[ind0].nonzero()[0]
    inde2 = inde2[ind0].nonzero()[0]
    p0 = [Poly0[:,:inde1[0]]]
    p1 = [Poly0[:,:inde1[0]]]
    for ii in range(0,len(inde1)):
        D0, D1 = Poly0[:,inde1[ii]],        Poly0[:,inde2[ii]]
        u0, u1 = D0-Poly0[:,inde1[ii]-1],   D1-Poly0[:,inde2[ii]+1]
        p0.append(_get_intersect(D0,u0,D1,u1)[:,np.newaxis])
        u0, u1 = Poly0[:,inde1[ii]+1]-D0,   Poly0[:,inde2[ii]-1]-D1
        p1.append(np.vstack([D0,_get_intersect(D0,u0,D1,u1),D1]).T)

    p0.append(Poly0[:,inde2[ii]+1:])
    p1.append(Poly0[:,inde2[ii]+1:])
    Poly0 = np.concatenate(tuple(p0),axis=1)
    Poly1 = np.concatenate(tuple(p1),axis=1)

    if save:
        nn = _name+'V0'
        pfe = os.path.join(path, nn+'.txt')
        np.savetxt(pfe, Poly0.T, comments='#',
                   header='Exp = %s\nName = %s\nCls = %s'%(_Exp,nn,_Cls))
        nn = _name+'V1'
        pfe = os.path.join(path, nn+'.txt')
        np.savetxt(pathfilext, Poly1.T, comments='#',
                   header='Exp = %s\nName = %s\nCls = %s'%(_Exp,nn,_Cls))
        nn = _name+'V2'
        pfe = os.path.join(path, nn+'.txt')
        np.savetxt(pathfilext, Poly.T, comments='#',
                   header='Exp = %s\nName = %s\nCls = %s'%(_Exp,nn,_Cls))
    return Poly0, Poly1, Poly, notes

if __name__=='__main__':

    # Parse input arguments
    msg = 'Launch creation of polygons txt from bash'
    parser = argparse.ArgumentParser(description = msg)

    parser.add_argument('-save', type=bool, help='save ?', default=_save)
    parser.add_argument('-path', type=str, help='saving path ?', default=_here)

    args = parser.parse_args()

    # Call wrapper function
    make_Poly(save=args.save, path=args.path)