#! /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 in ptypy (*.ptyd) format (in test !)

Example:
    pynx-ptycho-ptypy --data data.ptyd --probe focus,60e-6x200e-6,0.09 
      --algorithm analysis,ML**100,DM**200,nbprobe=3,probe=1 --verbose 10 
      --saveplot --liveplot
"""

params_beamline = {'xy': 'y,x'}

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


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

    def load_scan(self):
        h5data = h5py.File(self.params['data'], 'r')
        # TODO: handle multiple chunks
        tmp = h5data['/chunks/0/positions'][()]

        if True:
            self.x, self.y = tmp[:, 0], tmp[:, 1]
            do_transpose, do_flipud, do_fliplr = False, False, False

        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):
        imgn = self.imgn
        h5data = h5py.File(self.params['data'], 'r')
        self.params['instrument'] = self.params['data']
        self.params['nrj'] = h5data['/info/energy'][()]  # Energy in keV
        self.params['detectordistance'] = h5data['/info/distance'][()]

        self.params['pixelsize'] = h5data['/info/psize'][()]
        if np.isscalar(h5data['/info/rebin'][()]):
            self.params['pixelsize'] *= h5data['/info/rebin'][()]

        # Load all frames
        print("Reading frames from PtyPy file:")
        ii = 0
        t0 = timeit.default_timer()
        vimg = h5data['/chunks/0/data'][imgn.tolist()]

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

        self.raw_data = vimg
        self.load_data_post_process()


class PtychoRunnerPtyPy(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 = PtychoRunnerScanPtyPy

    @classmethod
    def make_parser(cls, default_par, description=None, script_name="pynx-ptycho-ptypy", epilog=None):
        if epilog is None:
            epilog = helptext_epilog
        if description is None:
            description = ("Script to perform a ptychography analysis on data recorded "
                           "using the PtyPy (.ptyd) format")
        p = default_par
        parser = super().make_parser(default_par, script_name, description, epilog)
        grp = parser.add_argument_group("PtyPy parameters")
        grp.add_argument('--data', type=str, default=None,
                         required=True, help='data file to analyse')
        # Change default value from main parser
        grp.add_argument('--xy', type=str, default=p['xy'],
                         help='expression to be used for the XY positions (e.g. "-x,y",...), to compute the'
                              'values from the raw input (scaling, flipping, etc..). Mathematical '
                              'operations can also be used, e.g.: ``--xy=0.5*x+0.732*y,0.732*x-0.5*y``')
        return parser


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