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

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


import sys
import os
import glob
import time
import locale
import timeit

import fabio
from ...utils import h5py
import numpy as np

from .runner import PtychoRunner, PtychoRunnerScan, PtychoRunnerException, default_params as params0
from .parser import ActionPtychoMotors

helptext_epilog = "  Examples:\n\n" \
                  "  * ``pynx-ptycho-id16a --h5meta meta.h5 --data data.h5 " \
                  "--probe 60e-6x60e-6,0.09 --ptychomotors=mot_pos.txt,-x,y " \
                  "--algorithm analysis,ML**100,DM**200,nbprobe=3,probe=1``\n\n" \
                  "  * ``pynx-ptycho-id16a --ptychomotors=mot_pos.txt,-x,y --flatfield=flat.npz " \
                  "--h5meta=meta.h5 --data=data.nxs --probe=focus,120e-6x120e-6,0.1 --defocus=500e-6 " \
                  "--algorithm=analysis,ML**100,DM**200,nbprobe=3,probe=1 " \
                  "--saveplot=object_phase --remove_obj_phase_ramp --liveplot``"

# NB: for id16 we start from a flat object (high energy, high transmission)
params_beamline = {
    'data': None,
    'h5meta': None,
    'imgname': None,
    'instrument': 'ESRF id16a',
    'monitor': None,
    'object': 'random,0.95,1,0,0.1',
    'pixelsize': 55e-6,
    'ptychomotors': None,
}

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


class PtychoRunnerScanID16a(PtychoRunnerScan):
    def __init__(self, params, scan, timings=None):
        super(PtychoRunnerScanID16a, self).__init__(params, scan, timings=timings)

    def load_scan(self):
        m = self.params['ptychomotors'].split(',')
        if len(m) in [1, 3]:
            print("reading motors from file:", m[0])
            # ptycho motor positions from single file, e.g. mot_pos.txt
            self.x, self.y = np.loadtxt(m[0], unpack=True)
            # Motor positions are in microns, convert to meters
            self.x *= 1e-6
            self.y *= 1e-6
        else:
            # ptycho motors from edf file
            xmot, ymot = m[0:2]

        if self.params['data'] is not None:
            imgn = np.arange(len(self.x), dtype=np.int32)
            if self.params['moduloframe'] is not None:
                n1, n2 = self.params['moduloframe']
                idx = np.where(imgn % n1 == n2)[0]
                imgn = imgn.take(idx)
                self.x = self.x.take(idx)
                self.y = self.y.take(idx)

            if self.params['maxframe'] is not None:
                N = self.params['maxframe']
                if len(imgn) > N:
                    print("MAXFRAME: only using first %d frames" % (N))
                    imgn = imgn[:N]
                    self.x = self.x[:N]
                    self.y = self.y[:N]
        else:
            # Load all available frames, as well as motor positions from frame headers (...)
            i = 0  # Start at 0 or 1 ?
            if self.params['imgname'].count('%') == 2:
                if os.path.isfile(self.params['imgname'] % (self.scan, i)) is False:
                    i = 1
            elif self.params['imgname'].count('%') == 3:
                if os.path.isfile(self.params['imgname'] % (self.scan, self.scan, i)) is False:
                    i = 1
            else:
                if os.path.isfile(self.params['imgname'] % i) is False:
                    i = 1

            vimg = []  # We don't know how many images there are, so start with a list...
            imgn = []
            x, y = [], []  # motor positions
            if self.params['monitor'] is not None:
                mon = []

            while True:
                if self.params['moduloframe'] is not None:
                    n1, n2 = self.params['moduloframe']
                    if i % n1 != n2:
                        i += 1
                        continue

                if i % 20 == 0:
                    sys.stdout.write('%d ' % i)
                    sys.stdout.flush()

                imgname = self.params['imgname']
                if imgname.count('%') == 2:
                    imgname = imgname % (self.scan, i)
                elif imgname.count('%') == 3:
                    imgname = imgname % (self.scan, self.scan, i)
                else:
                    imgname = imgname % i

                if os.path.isfile(imgname):
                    d = fabio.open(imgname)
                    vimg.append(d.data)
                    motors = {}
                    for k, v in zip(d.header['motor_mne'].split(), d.header['motor_pos'].split()):
                        motors[k] = float(v)
                    x.append(motors[xmot])  # Motor positions in microns
                    y.append(motors[ymot])
                    imgn.append(i)
                    if self.params['monitor'] is not None:
                        counters = {}
                        for k, v in zip(d.header['counter_mne'].split(), d.header['counter_pos'].split()):
                            counters[k] = float(v)
                        mon.append(counters[self.params['monitor']])
                    if self.params['maxframe'] is not None:
                        if len(vimg) == self.params['maxframe']:
                            print("MAXFRAME: only using first %d frames" % (len(vimg)))
                            break
                    i += 1
                else:
                    break
            self.x = np.array(x) * 1e-6
            self.y = np.array(y) * 1e-6
            imgn = np.array(imgn)

        if len(m) >= 3:
            x, y = self.x, self.y
            self.x, self.y = eval(m[-2]), eval(m[-1])

        if len(self.x) < 4:
            raise PtychoRunnerException("Less than 4 scan positions, is this a ptycho scan ?")

        if self.params['monitor'] is not None:
            self.raw_data_monitor = np.array(mon)
            mon0 = np.median(mon)
            mon /= mon0
            self.validframes = np.where(mon > 0.1)
            if len(self.validframes) != len(mon):
                print('WARNING: The following frames have a monitor value < 0.1 x '
                      'the median value and will be ignored (no beam ?)')
                print(np.where(mon <= (mon0 * 0.1)))
            self.x = np.take(self.x, self.validframes)
            self.y = np.take(self.y, self.validframes)
            imgn = np.take(imgn, self.validframes)
            self.raw_data_monitor = np.take(self.raw_data_monitor, self.validframes)

        self.imgn = imgn

    def load_data(self):
        h5meta_filename = self.params['h5meta']
        if '%' in h5meta_filename:
            h5meta_filename = h5meta_filename % tuple(self.scan for i in range(h5meta_filename.count('%')))
            print('h5meta filename for scan #%d: %s' % (self.scan, h5meta_filename))
        h5meta = h5py.File(h5meta_filename, 'r')

        # Assume there's only a single entry, of which we don't know the name...
        for v in h5meta.values():
            break

        h5prefix = v.name
        # DEBUG
        self.h5prefix = h5prefix
        self.h5meta = h5meta

        date_string = h5meta[h5prefix + "/start_time"][()]  # '2015-03-11T23:56:58.390629'
        if sys.version_info > (3,) and isinstance(date_string, bytes):
            date_string = date_string.decode('utf-8')  # Could also use ASCII in this case
        else:
            date_string = str(date_string)
        pattern = '%Y-%m-%dT%H:%M:%S.%f'
        try:
            lc = locale._setlocale(locale.LC_ALL)
            locale._setlocale(locale.LC_ALL, 'C')
            epoch = int(time.mktime(time.strptime(date_string, pattern)))
            locale._setlocale(locale.LC_ALL, lc)
        except ValueError:
            print("Could not extract time from spec header, unrecognized format: %s, expected: %s" % (
                date_string, pattern))

        self.params['nrj'] = h5meta.get(h5prefix + "/measurement/initial/energy")[()]
        if self.scan is None:
            self.scan = int(h5meta.get(h5prefix + "/scan_number")[()])
        if self.params['detectordistance'] is None:
            # self.params['detectordistance'] = h5meta.get(h5prefix+"/PTYCHO/focusToDetectorDistance")[()]
            tmp = eval(h5meta.get(h5prefix + "/PTYCHO/parameters")[()])
            self.params['detectordistance'] = float(tmp['focusToDetectorDistance'])  # Seriously ??
            print("Detector distance from h5 metadata: %6.3fm" % self.params['detectordistance'])

        # read all frames
        imgn = self.imgn
        t0 = timeit.default_timer()
        sys.stdout.write("Reading frames:")
        sys.stdout.flush()

        # Load all images from a single hdf5/nexus file
        # Here assuming lambda detector pseudo-nexus files
        h5data_filename = self.params['data']
        if h5data_filename.count('%') == 1:
            h5data_filename = h5data_filename % tuple(self.scan for i in range(h5data_filename.count('%')))
        print('data filename for scan #%d: %s' % (self.scan, h5data_filename))
        with h5py.File(h5data_filename, 'r') as h:
            data = h["/entry/instrument/detector/data"]
            vimg = data[imgn.tolist()]
            # mask ?
            if '/entry/instrument/detector/pixel_mask' in h:
                mask = h['/entry/instrument/detector/pixel_mask']
                shape_mismatch = False
                if mask.ndim == 3:
                    if mask.shape[0] == 1:
                        # Same mask for all frames
                        mask = mask[0]
                    elif mask.shape[0] == h["/entry/instrument/detector/data"].shape[0]:
                        # Mask individualised per frame
                        mask = mask[imgn.tolist()]
                    else:
                        shape_mismatch = True
                if mask.shape[-2:] != data.shape[-2:]:
                    shape_mismatch = True
                if shape_mismatch:
                    print(f"WARNING: the shape of the pixel mask [{h5data_filename}:"
                          f"/entry/instrument/detector/pixel_mask: {mask.shape}] "
                          f"does not match the detector data shape: {data.shape}- "
                          f"IGNORING this mask")
                else:
                    nb = np.sum(mask)
                    if nb:
                        self.raw_mask = mask.astype(np.int8)
                        print(f"Loading detector mask from {h5data_filename}:"
                              f"/entry/instrument/detector/pixel_mask with "
                              f"{nb} [{nb / mask.size * 100:.2f}%] pixels masked")

        dt = timeit.default_timer() - t0
        print("\n")
        print('Time to read all frames: %4.1fs [%5.2f Mpixel/s]' % (dt, vimg[0].size * len(vimg) / 1e6 / dt))

        # Convert to numpy arrays
        d = np.empty((len(vimg), vimg[0].shape[0], vimg[0].shape[1]), dtype=vimg[0].dtype)
        for i in range(len(vimg)):
            d[i] = vimg[i]

        self.raw_data = d
        self.load_data_post_process()


class PtychoRunnerID16a(PtychoRunner):
    """
    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(PtychoRunnerID16a, self).__init__(argv, default_params if params is None else params)
        self.PtychoRunnerScan = PtychoRunnerScanID16a

    @classmethod
    def make_parser(cls, default_par, description=None, script_name="pynx-ptycho-id16a", epilog=None):
        if epilog is None:
            epilog = helptext_epilog
        if description is None:
            description = "Script to perform a ptychography analysis on FAR FIELD data from ID16A@ESRF"
        p = default_par

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

        grp.add_argument('--h5meta', '--meta', type=str, default=None, required=True,
                         help='hdf5 filename with the scan metadata.'
                              'The energy and detector distance will be extracted from it')

        grp.add_argument('--data', '--h5data', type=str, default=None, required=True,
                         help='path to the hdf5 file with all frames, e.g. /some/dir/to/data.h5')

        grp.add_argument('--ptycho_motors', '--ptychomotors', type=str, required=True,
                         dest='ptychomotors',
                         help="name of the two motors used for ptychography, or a motor "
                              "positions file, optionally followed by a mathematical expression to"
                              "be used to calculate the actual motor positions (axis convention, angle..)."
                              "Values are either:\n\n"
                              "* extracted from the data files, and are assumed to be in mm."
                              " [this is currently not available]\n"
                              "* or extracted from the 'mot_pos.txt' file, and assumed to"
                              "be in microns..\n"
                              "Examples:\n\n"
                              "* ``--ptychomotors=mot_pos.txt,-x,y``\n"
                              "* ``--ptychomotors spy spz``\n"
                              "* ``--ptychomotors=spy,spz,-x,y``\n"
                              "* ``--ptychomotors mot_pos.txt``\n"
                              "Note that if the ``xy=-y,x`` command-line argument is used, it is "
                              "applied after this, using ``--ptychomotors=mot_pos.txt,-x,y`` "
                              "is equivalent to ``--ptychomotors mot_pos.txt --xy=-x,y``")

        grp.add_argument('--monitor', type=str, default=None,
                         help='optional name for the monitor counter. The frames will '
                              'be normalized by the ratio of the counter value divided '
                              'by the median value of the counter over the entire'
                              'scan (so as to remain close to Poisson statistics). '
                              'A monitor intensity lower than 10%% of the median value '
                              'will be interpreted as an image taken without beam'
                              'and will be skipped.')
        return parser


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