#! /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 warnings
import argparse

from ...utils import h5py
import numpy as np
from silx.io.dictdump import nxtodict

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

helptext_epilog = """
Examples:

* ``pynx-ptycho-id16a-nf --data data.h5 --meta meta.h5 --dark dark000.h5\
--algorithm analysis,ML**100,DM**200,nbprobe=3,probe=1``

Using a reference (direct beam) frame, allows to use Paganin, and 
with use_direct_beam will provide an absolute scale for the probe.
Also, only save the object phase for a smaller output:
    
* ``pynx-ptycho-id16a-nf --data data.h5 --meta meta.h5 --dark dark000.h5 --data_ref ref.h5\
--algorithm analysis,ML**100,DM**200,Paganin,probe=1\
--use_direct_beam --cxi_output object_phase``

"""

# NB: for id16 we start from a flat object (high energy, high transmission)
params_beamline = {
    'adu_scale': None,
    'data_ref': None,
    'delta_beta': None,
    'autocenter': False,
    'instrument': 'ESRF id16a',
    'maxsize': 10000,
    'meta': None,
    'near_field': True,
    'object': 'random,0.95,1,0,0.1',
    'padding': 0,
    'ptychomotors': ['spy', 'spz', '-x*1e-6', 'y*1e-6'],
    'roi': 'full',
    'tomo_motor': 'somega'
}

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


class PtychoRunnerScanID16aNF(PtychoRunnerScan):
    def __init__(self, params, scan, timings=None):
        super(PtychoRunnerScanID16aNF, self).__init__(params, scan, timings=timings)
        self.h5_header_path = None
        self.tomo_metadata = {}

    def load_scan(self):
        data_filename = self.params['data']
        if '%' in data_filename:
            data_filename = data_filename % tuple(self.scan for i in range(data_filename.count('%')))
            print('data filename for scan #%d: %s' % (self.scan, data_filename))
        data = h5py.File(data_filename, mode='r', enable_file_locking=False)
        m = self.params['ptychomotors']
        xmot, ymot = m[0:2]
        for instrument in ['instrument', 'ESRF-ID16A']:
            for detector_name in ['Frelon', 'PCIe', 'ximea']:
                s = f'entry_0000/{instrument}/{detector_name}/header'
                if s in data:
                    self.h5_header_path = s
                    break
            if self.h5_header_path is not None:
                break
        self.x = np.array([v for v in data[f'{self.h5_header_path}/{xmot}'][()].split()], dtype=np.float32)
        self.y = np.array([v for v in data[f'{self.h5_header_path}/{ymot}'][()].split()], dtype=np.float32)
        # Extract tomo angle
        motors = {}
        header = data[f'{self.h5_header_path}']

        for k, v in zip(header['motor_mne '][()].split(), header['motor_pos '][()].split()):
            motors[k.decode()] = float(v)
        ang = np.deg2rad(motors[self.params['tomo_motor']])
        self.tomo_metadata['angle'] = ang
        print(f"Tomography angle: {self.params['tomo_motor']}={np.rad2deg(ang):8.3f}")
        self.tomo_metadata['data'] = data_filename

        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]

        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 ?")

        self.imgn = imgn

    def load_data(self):
        data_filename = self.params['data']
        if '%' in data_filename:
            data_filename = data_filename % tuple(self.scan for i in range(data_filename.count('%')))
        data = h5py.File(data_filename, mode='r', enable_file_locking=False)

        meta = h5py.File(self.params['meta'], mode='r', enable_file_locking=False)
        # Entry name is obfuscated, so just take the first entry with 'entry' in the name
        meta_entry = None
        for k, v in meta.items():
            if 'entry' in k:
                meta_entry = v
                break

        # NXsample
        if 'sample' in meta_entry:
            try:
                self.sample_nx_dict = nxtodict(meta_entry['sample'])
                if 'name' in self.sample_nx_dict:
                    self.sample_name = self.sample_nx_dict['name'].tobytes().decode()
            except Exception:
                warnings.warn("Something went wrong converting NXsample info - skipping")

        if False:
            date_string = data["entry_0000/start_time"][()]  # '2020-12-12T15:29:07Z'
            if 'Z' == date_string[-1]:
                date_string = date_string[:-1]
                # TODO: for python 3.7+, use datetime.isoformat()
            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'

            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))

        if self.params['nrj'] is None:
            self.params['nrj'] = float(meta_entry["instrument/monochromator/energy"][()])
            print("Energy read from the meta file: %6.3f keV" % self.params['nrj'])
        else:
            print("Energy from the command line: %6.3f keV" % self.params['nrj'])

        if self.params['pixelsize'] is None:
            self.params['pixelsize'] = float(meta_entry["TOMO/pixelSize"][()]) * 1e-6
            print("Effective pixel size: %12.3fnm [from meta file]" % (self.params['pixelsize'] * 1e9))
        else:
            print("Effective pixel size: %12.3fnm [from command-line]" % (self.params['pixelsize'] * 1e9))

        if self.scan is None:
            self.scan = 0

        # Raw positioners positions, almost guaranteed not to be SI
        positioners = {}
        tmpn = meta_entry["sample/positioners/name"][()]
        tmpv = meta_entry["sample/positioners/value"][()]
        if isinstance(tmpn, np.bytes_):
            tmpn = tmpn.decode('ascii')
            tmpv = tmpv.decode('ascii')

        for k, v in zip(tmpn.split(), tmpv.split()):
            positioners[k] = float(v)

        if self.params['detectordistance'] is None:
            if f'{self.h5_header_path}/propagation_distance' in data:
                self.params['detectordistance'] = \
                    1e-3 * float(data[f'{self.h5_header_path}/propagation_distance'][()])
                print(f"Effective propagation distance (from {self.h5_header_path}/propagation_distance): "
                      "%12.8fm" % self.params['detectordistance'])
            else:
                sx = 1e-3 * positioners['sx']
                sx0 = 1e-3 * float(meta_entry["TOMO/sx0"][()])
                z1 = sx - sx0
                z12 = 1e-3 * float(meta_entry["PTYCHO/focusToDetectorDistance"][()])  # z1+z2
                z2 = z12 - z1
                print(sx * 1e3, sx0 * 1e3, z1 * 1e3, z12 * 1e3, z2 * 1e3)
                self.params['detectordistance'] = z1 * z2 / z12
                print("Effective propagation distance (computed): %12.8fm" % self.params['detectordistance'])

        # read all frames
        imgn = self.imgn
        t0 = timeit.default_timer()
        print("Reading frames:")
        self.raw_data = data["/entry_0000/measurement/data"][imgn.tolist()]
        dt = timeit.default_timer() - t0
        print('Time to read all frames: %4.1fs [%5.2f Mpixel/s]' %
              (dt, self.raw_data[0].size * len(self.raw_data) / 1e6 / dt))

        if self.params['adu_scale'] is not None:
            self.raw_data = self.raw_data.astype(np.float32, copy=False)
            self.raw_data *= self.params['adu_scale']
        # Read empty beam / reference images if necessary
        if self.params['data_ref'] is not None:
            # Assume this is an hdf5 file with a standard path for now
            with h5py.File(self.params['data_ref']) as h:
                d = h['entry_0000/measurement/data'][()].astype(np.float32, copy=False)
            if d.ndim == 3:
                self.data_ref = d.mean(axis=0)
            else:
                self.data_ref = d
            if self.params['adu_scale'] is not None:
                self.data_ref = self.data_ref.astype(np.float32, copy=False)
                self.data_ref *= self.params['adu_scale']

        self.load_data_post_process()


class PtychoRunnerID16aNF(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(PtychoRunnerID16aNF, self).__init__(argv, default_params if params is None else params)
        self.PtychoRunnerScan = PtychoRunnerScanID16aNF

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

        parser = super().make_parser(p, script_name, description, epilog)
        grp = parser.add_argument_group("ID16A (near field) parameters")
        grp.add_argument('--data', type=str, default=p['data'], required=True,
                         help="path to the bliss file, e.g. path/to/data.h5."
                              "This can be a pattern, e.g. --data 'sample1_%%04d.h5', in which "
                              "case the field will be replaced by the scan number.")

        grp.add_argument('--meta', type=str, default=p['meta'], required=True,
                         help='path to the bliss metadata file, e.g. path/to/meta.h5.')

        grp.add_argument('--data_ref', type=str, default=p['data_ref'],
                         help="path to hdf5 file with reference images (a.k.a. 'empty beam')"
                              "of the direct beam, which are needed for Paganin and CTF algorithms, "
                              "e.g. '--data_ref path/to/ref.h5'")

        grp.add_argument('--ptycho_motors', '--ptychomotors', type=str, default=p['ptychomotors'],
                         nargs='+', dest='ptychomotors', action=ActionPtychoMotors,
                         help="name of the two motors used for ptychography, optionally followed "
                              "by a mathematical expression to be used to calculate the actual "
                              "motor positions (axis convention, angle..) in meters. Values will be read "
                              "from the dat files:\n\n"
                              "* ``--ptychomotors=spy,spz,-x*1e-6,y*1e-6``  (the default)\n"
                              "* ``--ptychomotors pix piz``\n"
                              "Note that if the ``--xy=-y,x`` command-line argument is used, "
                              "it is applied _after_ this, using ``--ptychomotors=spy,spz,-x,y`` "
                              "is equivalent to ``--ptychomotors spy spz --xy=-x,y``")

        grp.add_argument('--delta_beta', type=float, default=p['delta_beta'],
                         help="delta/beta value, required for Paganin or CTF algorithms."
                              "Can also be set in the algorithm string")

        grp.add_argument('--adu_scale', type=float, default=p['adu_scale'],
                         help="optional scale factor by which the all data (including data_ref and dark) "
                              "will be multiplied in order to have photon counts.")

        grp.add_argument('--tomo_motor', type=str, default=p['tomo_motor'],
                         help='Name of the tomography motor -only used to export the rotation '
                              'angle in output files for ptycho-tomo.')

        grp.add_argument('--padding', type=int, default=p['padding'],
                         help="padding value [DEPRECATED, does not work]")

        # Change const value
        grp.add_argument('--save_plot', '--saveplot', type=str,
                         default=False, const='object_phase', dest='saveplot',
                         choices=['object_phase', 'object_rgba'], nargs='?',
                         help='will save plot at the end of the optimization (png file).\n'
                              'A string can be given to specify if only the object phase '
                              '(default if only --saveplot is given) or rgba should be plotted.')

        # Experimental
        # grp.add_argument('--prefilter', type=int, default=0, nargs='?', const=200,
        #                  help=argparse.SUPPRESS)
        # "If --data_ref is supplied, it is possible to normalise "
        #      "all images by a low-passed filtered direct beam, "
        #      "in order to avoid intensity discontinuity at the periodic "
        #      "image boundaries.")
        grp.add_argument('--tapering', type=int, default=0, nargs='?', const=200,
                         help=argparse.SUPPRESS)
        # "Add a tapering (cosine/Tukey) window so that the intensities
        # "reach zero on the boundaries of the observed intensity data")`

        return parser

    def check_params_beamline(self):
        """
        Check if self.params includes a minimal set of valid parameters, specific to a beamline
        Returns: Nothing. Will raise an exception if necessary
        """
        if 'paganin' in self.params['algorithm'].lower() or 'ctf' in self.params['algorithm'].lower():
            if self.params['delta_beta'] is None and 'delta_beta:' not in self.params['algorithm'].lower():
                raise PtychoRunnerException('Need delta_beta=... for Paganin and CTF')
            if self.params['data_ref'] is None:
                raise PtychoRunnerException('Need data_ref=.../ref.h5 for Paganin or CTF')
        if self.params['dark'] is not None:
            if self.params['dark'].find('.h5') > 0 or self.params['dark'].find('.hdf5') > 0:
                if self.params['dark'][-3:] == '.h5' or self.params['dark'][-5:] == '.hdf5' \
                        and ":" not in self.params['dark']:
                    self.params['dark'] += ':/entry_0000/measurement/data'
                    print("dark: hdf5 path missing, adding it => dark=%s" % self.params['dark'])


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