#! /opt/local/bin/python
# -*- coding: utf-8 -*-

# PyNX - Python tools for Nano-structures Crystallography
#   (c) 2017-present : ESRF-European Synchrotron Radiation Facility
#       authors:
#         Vincent Favre-Nicolin, favre@esrf.fr

import sys
import timeit
import numpy as np
import fabio
import h5py
import os
from silx.io.specfile import SpecFile
from .runner import CDIRunner, CDIRunnerException, CDIRunnerScan, default_params as params0
from pynx.cdi import *

helptext_epilog = " \n\nExamples:\n" \
                  " * ``pynx-cdi-id01 --data alignment_S2280.cxi --liveplot``: just using default parameters\n" \
                  " * ``pynx-cdi-id01 --data alignment_S2280.cxi --liveplot --support_threshold 0.2``:\n" \
                  "   changing the support threshold\n" \
                  " * ``pynx-cdi-id01 --data alignment_S2280.cxi --liveplot --support_threshold 0.2 0.3 " \
                  "--nb_run 10 --nb_run_keep 5``:\n" \
                  "   Generate 10 solutions with random support threshold between 0.2 and 0.3, and keep" \
                  "   only the 5 best."

params_beamline = {
    'detector': None,
    'detwin': True,
    'imgcounter': 'auto',
    'imgname': None,
    'instrument': 'ESRF id01',
    'nb_er': 200,
    'nb_hio': 0,
    'nb_ml': 0,
    'nb_raar': 600,
    'roi': 'auto',
    'scan': None,
    'specfile': None,
    'support_size': None,
    'zero_mask': 'auto',
}

default_params = params0.copy()
for k, v in params_beamline.items():
    default_params[k] = v


class CDIRunnerScanID01(CDIRunnerScan):
    def __init__(self, params, scan, timings=None):
        super(CDIRunnerScanID01, self).__init__(params, scan, timings=timings)

    def load_data(self):
        """
        Loads data. If no id01-specific keywords have been supplied, use the default data loading.

        """
        if self.params['specfile'] is not None:
            # SPEC
            if self.scan is None:
                raise CDIRunnerException("Must supply scan=N along with specfile")
            if isinstance(self.scan, str):
                # do we combine several scans ?
                vs = self.scan.split('+')
            else:
                vs = [self.scan]
            imgn = None
            scan_motor_last_value = None
            for scan in vs:
                if scan is None:
                    scan = 0
                else:
                    scan = int(scan)
                s = SpecFile(self.params['specfile'])['%d.1' % (scan)]
                h = s.scan_header_dict

                if self.params['imgcounter'] == 'auto':
                    if 'ei2minr' in s.labels:
                        self.params['imgcounter'] = 'ei2minr'
                    elif 'mpx4inr' in s.labels:
                        self.params['imgcounter'] = 'mpx4inr'
                    print("Using image counter: %s" % (self.params['imgcounter']))

                if self.params['wavelength'] is None and 'UMONO' in h:
                    nrj = float(h['UMONO'].split('mononrj=')[1].split('ke')[0])
                    w = 12.3984 / nrj
                    self.params['wavelength'] = w
                    print("Reading nrj from spec data: nrj=%6.3fkeV, wavelength=%6.3fA" % (nrj, w))

                if 'UDETCALIB' in h:
                    # UDETCALIB cen_pix_x=18.347,cen_pix_y=278.971,pixperdeg=445.001,det_distance_CC=1402.175,
                    # det_distance_COM=1401.096,timestamp=20170926...
                    if self.params['detector_distance'] is None:
                        if 'det_distance_CC=' in h['UDETCALIB']:
                            self.params['detector_distance'] = float(
                                h['UDETCALIB'].split('stance_CC=')[1].split(',')[0])
                            print("Reading detector distance from spec data: %6.3fm" % self.params[
                                'detector_distance'])
                        else:
                            print(
                                'No detector distance given. No det_distance_CC in UDETCALIB ??: %s' % (
                                    h['UDETCALIB']))
                # Also save specfile parameters as a 'user_' param
                self.params['user_id01_specfile_header'] = s.file_header_dict
                self.params['user_id01_specfile_scan_header'] = s.scan_header_dict

                # Read images
                if imgn is None:
                    imgn = s.data_column_by_name(self.params['imgcounter']).astype(int)
                else:
                    # try to be smart: exclude first image if first motor position is at the end of last scan
                    if scan_motor_last_value == s.data[0, 0]:
                        i0 = 0
                    else:
                        print("Scan %d: excluding first image at same position as previous one" % (scan))
                        i0 = 1
                    imgn = np.append(imgn, s.data_column_by_name(self.params['imgcounter'])[i0:].astype(int))
                scan_motor_last_value = s.data[0, 0]

                self.iobs = None
                t0 = timeit.default_timer()
                if self.params['imgname'] is None:
                    if 'ULIMA_eiger2M' in h:
                        imgname = h['ULIMA_eiger2M'].strip().split('_eiger2M_')[0] + '_eiger2M_%05d.edf.gz'
                        print("Using Eiger 2M detector images: %s" % (imgname))
                    else:
                        imgname = h['ULIMA_mpx4'].strip().split('_mpx4_')[0] + '_mpx4_%05d.edf.gz'
                        print("Using Maxipix mpx4 detector images: %s" % (imgname))
                else:
                    imgname = self.params['imgname']
                sys.stdout.write('Reading frames: ')
                ii = 0
                for i in imgn:
                    if (i - imgn[0]) % 20 == 0:
                        sys.stdout.write('%d ' % (i - imgn[0]))
                        sys.stdout.flush()
                    frame = fabio.open(imgname % i).data
                    if self.iobs is None:
                        self.iobs = np.empty((len(imgn), frame.shape[0], frame.shape[1]), dtype=frame.dtype)
                    self.iobs[ii] = frame
                    ii += 1
                print("\n")
                dt = timeit.default_timer() - t0
                print('Time to read all frames: %4.1fs [%5.2f Mpixel/s]' % (dt, self.iobs.size / 1e6 / dt))

        elif self.params['data'] is not None and self.scan is not None:
            # BLISS file ?
            scan = int(self.scan)
            bliss_data = True
            try:
                with h5py.File(self.params['data'], 'r') as h5file:
                    if h5file[f'/{scan}.1/instrument/title'][()].decode("latin-1") != 'ESRF: id01':
                        print(h5file[f'/{scan}.1/instrument/title'][()])
                        bliss_data = False
            except KeyError:
                bliss_data = False
            except OSError:
                bliss_data = False

            if bliss_data:
                if os.path.isfile(self.params['data']) is False:
                    raise CDIRunnerException("hdf5 file does not exist: %s" % (self.params['data']))
                print("Reading BLISS NeXus data file")
                h5file = h5py.File(self.params['data'], 'r')
                self.params['instrument'] = h5file[f'/{self.scan}.1/instrument/title'][()]
                if self.params['wavelength'] is None:
                    self.params['wavelength'] = h5file[f'/{self.scan}.1/instrument/monochromator/WaveLength'][()]
                    # / 1.60218e-16

                det = self.params['detector']
                if det is None:
                    # Identify detector
                    smpx1x4 = f"/{self.scan}.1/measurement/mpx1x4"
                    seiger2M = f"/{self.scan}.1/measurement/eiger2M"
                    if smpx1x4 in h5file:
                        det = "mpx1x4"
                        if seiger2M in h5file:
                            raise CDIRunnerException("You need to select either detector=mpx1x4 or eiger2M "
                                                     "- both are present")
                    elif seiger2M in h5file:
                        det = "eiger2M"
                    else:
                        raise CDIRunnerException(
                            "You need to supply detector=... ; neither mpx1x4 or eiger2M was found")
                    self.params['detector'] = det

                if self.params['detector_distance'] is None and \
                        f"/{self.scan}.1/instrument/{det}/distance" in h5file:
                    self.params['detector_distance'] = \
                        h5file[f"/{self.scan}.1/instrument/{det}/distance"][()]
                    print("Detector distance:", self.params['detector_distance'])

                if self.params['pixel_size_detector'] is None and \
                        f"/{self.scan}.1/instrument/{det}/x_pixel_size" in h5file:
                    self.params['pixel_size_detector'] = \
                        h5file[f"/{self.scan}.1/instrument/{det}/x_pixel_size"][()]
                    print("Detector pixel size:", self.params['pixel_size_detector'])

                # Read images
                x = h5file[f"{self.scan}.1/measurement/{det}"].shape[0]
                imgn = np.arange(x, dtype=int)
                t0 = timeit.default_timer()
                sys.stdout.write('Reading frames: ')
                ii = 0
                for i in imgn:
                    if (i - imgn[0]) % 20 == 0:
                        sys.stdout.write('%d ' % (i - imgn[0]))
                        sys.stdout.flush()
                    frame = h5file[f"/{self.scan}.1/measurement/{self.params['detector']}"][i, :, :]
                    if self.iobs is None:
                        self.iobs = np.empty((len(imgn), frame.shape[0], frame.shape[1]), dtype=frame.dtype)
                    self.iobs[ii] = frame
                    ii += 1
                print("\n")
                dt = timeit.default_timer() - t0
                print('Time to read all frames: %4.1fs [%5.2f Mpixel/s]' % (dt, self.iobs.size / 1e6 / dt))

                if f"/{self.scan}.1/instrument/{det}/pixel_mask" in h5file:
                    mask = h5file[f"/{self.scan}.1/instrument/{det}/pixel_mask"][()]
                    if np.ndim(mask) >= 2:
                        self.mask = (mask != 0).astype(np.int8)
                        print(f"Loaded mask from bliss nexus data: /{self.scan}.1/instrument/{det}/pixel_mask")
            else:
                # Likely some other generic data format (hdf5, cxi, tiff, edf...)
                super(CDIRunnerScanID01, self).load_data()
        else:
            # Likely some other generic data format (hdf5, cxi, tiff, edf...)
            super(CDIRunnerScanID01, self).load_data()

    def prepare_cdi(self):
        """
        Prepare CDI object from input data.

        :return: nothing. Creates or updates self.cdi object.
        """
        super(CDIRunnerScanID01, self).prepare_cdi()
        # Scale initial object (unnecessary if auto-correlation is used)
        if self.params['support'] != 'auto':
            self.cdi = ScaleObj(method='F', lazy=True) * self.cdi


class CDIRunnerID01(CDIRunner):
    """
    Class to process a series of scans with a series of algorithms, given from the command-line
    """

    def __init__(self, argv, params, *args, **kwargs):
        super().__init__(argv, params)
        self.CDIRunnerScan = CDIRunnerScanID01

    @classmethod
    def make_parser(cls, description=None, script_name="pynx-ptycho-id01", epilog=None, default_par=None):
        if default_par is None:
            default_par = default_params
        if epilog is None:
            epilog = helptext_epilog
        if description is None:
            description = "Script to perform a CDI analysis on data from ID01@ESRF\n\n" \
                          "See examples at the end."

        parser = super().make_parser(script_name, description, epilog, default_par)
        grp = parser.add_argument_group("ID01 parameters")

        grp.add_argument('--detector', type=str, default=None,
                         help="detector used, e.g. either 'mpx1x4' or 'eiger2M'.\n"
                              "This is normally auto-detected if only one detector data is present.")

        grp.add_argument(
            '--scan', type=str, default=None,
            help="scan number in bliss (or spec) file"
                 "Alternatively, a list or range of scans can be given.\n"
                 "Examples:\n"
                 "  --scan 11\n"
                 "  --scan 12,23,45\n"
                 "  --scan 'range(12,25)' (note the quotes).")

        grp = parser.add_argument_group("OBSOLETE ID01 parameters")

        grp.add_argument('--specfile', type=str, default=None,
                         help='/some/dir/to/specfile.spec: path to specfile [obsolete]')

        grp.add_argument('--imgcounter', type=str, default=None, choices=['mpx4inr', 'ei2mint'],
                         help='spec counter name for image number [only for SPEC data-obsolete]')

        grp.add_argument('--imgname', type=str, default=None,
                         help="/dir/to/images/prefix%%05d.edf.gz: images location with prefix (for SPEC data)"
                              "[default: will be extracted from the ULIMA_mpx4 entry in the spec scan header]")
        return parser


def make_parser_sphinx():
    """Returns the argparse for sphinx documentation"""
    return CDIRunnerID01.make_parser()
