#! /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 timeit

from ...utils import h5py
import numpy as np

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

helptext_epilog = """
Script to perform a ptychography analysis on data recorded on nanoscopium@Soleil (*.nxs format)

Example:
    pynx-ptycho-nanoscopium --data flyscan_1061-0001.nxs --probe focus,700e-6,0.63
        --defocus 4 --detectordistance 3.265
        --algorithm analysis,ML**100,AP**400,DM**100,nprobe=3,DM**200,probe=1 
        --saveplot --mask maxipix --liveplot

"""

params_beamline = {'instrument': 'nanoscopium@Soleil', 'pixelsize': 55e-6, 'mask': 'maxipix',
                   'object': 'random,0.95,1,0,0.1', 'obj_smooth': 0.5, 'obj_inertia': 0.1, 'probe_smooth': 0.2,
                   'probe_inertia': 0.01}

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


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

    def load_scan(self):
        filename = self.params['data']
        if '%' in filename:
            filename = filename % self.scan
            self.print('Loading data:', filename)
        h5data = h5py.File(filename, 'r')
        # Get the first entry and hope for the best
        entry = [v for v in h5data.values()][0]

        if 'scan_data/sample_Piezo_TX' in entry:
            self.x = entry['scan_data/sample_Piezo_TX'][()].flatten() * -1e-3
        else:
            self.x = entry['scan_data/sample_Piezo_Tx'][()].flatten() * -1e-3
        if 'scan_data/sample_Piezo_Tz' in entry:
            self.y = entry['scan_data/sample_Piezo_Tz'][()].flatten() * 1e-3
        else:
            self.y = entry['scan_data/sample_Piezo_TZ'][()].flatten() * 1e-3

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

        imgn = np.arange(len(self.x), dtype=int)

        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]
        self.imgn = imgn

    def load_data(self):
        filename = self.params['data']
        if '%' in filename:
            filename = filename % self.scan
        h5data = h5py.File(filename, 'r')
        # Get the first entry and hope for the best
        entry = [v for v in h5data.values()][0]
        if self.params['nrj'] is None:
            # Energy in keV must be corrected (ask Kadda !)
            self.params['nrj'] = entry['NANOSCOPIUM/Monochromator/energy'][0] * 1.2371 - 1.3456

        # Load all frames
        print("Reading frames from nxs file:")
        t0 = timeit.default_timer()
        h5d = entry['scan_data/Image_merlin_image']
        vimg = None
        if len(h5d.shape) == 4:
            imgn = np.array(self.imgn, dtype=np.int32)
            if vimg is None:
                vimg = np.empty((self.x.size, h5d.shape[-2], h5d.shape[-1]))
            # flyscan data is 2D..
            n1, n2 = h5d.shape[:2]
            for i in range(n1):
                i0 = n2 * i
                idx = np.where(np.logical_and(imgn >= i0, imgn < (i0 + n2)))[0]
                if len(idx):
                    vimg[idx] = h5d[i, imgn[idx] - i0]
        else:
            vimg = h5d[self.imgn]

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

        # Add gaps
        n, ny, nx = vimg.shape
        self.raw_data = np.empty((n, 516, 516))
        self.raw_data[:, :256, :256] = vimg[:, :256, :256]
        self.raw_data[:, 260:, :256] = vimg[:, 256:, :256]
        self.raw_data[:, :256, 260:] = vimg[:, :256, 256:]
        self.raw_data[:, 260:, 260:] = vimg[:, 256:, 256:]
        # TODO: check the following values - pixels should be masked anyway...
        self.raw_data[:, 256] *= 2.5
        self.raw_data[:, :, 256] *= 2.5
        self.raw_data[:, 260] *= 2.5
        self.raw_data[:, :, 260] *= 2.5

        print(self.x.shape, self.y.shape, self.raw_data.shape)

        self.load_data_post_process()


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

    @classmethod
    def make_parser(cls, default_par, description=None, script_name="pynx-ptycho-nanoscopium", epilog=None):
        if epilog is None:
            epilog = helptext_epilog
        if description is None:
            description = ("Script to perform a ptychography analysis on data recorded "
                           "in CXI format (http://cxidb.org/cxi.html)")

        parser = super().make_parser(default_par, script_name, description, epilog)
        p = default_par
        grp = parser.add_argument_group("Nanoscopium@Soleil parameters")
        grp.add_argument('--data', type=str, default=None,
                         required=True, help='data file (hdf5) to analyse')
        # Required for this runner (not in data file)
        # This overrides the main parser
        grp.add_argument('--nrj', '--energy', action='store', default=p['nrj'],
                         type=float, required=True, help='X-ray beam energy in keV')
        grp.add_argument('--detector_distance', '--detectordistance', '--distance',
                         type=float, dest='detectordistance', default=p['detectordistance'],
                         required=True, help='Detector distance in meters.')

        return parser


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