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

from ...utils import h5py
import numpy as np

import silx

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

helptext_epilog = """
Script to perform a ptychography analysis on data from cristal@Soleil

Example:
    pynx-ptycho-cristal --data my_nexus_file.nxs --h5data /a/scan_data/data_04 
      --ptychomotors /a/scan_data/trajectory_1_1 /a/scan_data/trajectory_1_2 
      --detectordistance 1.3 --nrj 8 --xy x*1e-12,y*1e-12 --probe 60e-6x200e-6,0.09 
      --algorithm analysis,ML**100,DM**200,nbprobe=3,probe=1
      --verbose 10 --saveplot --liveplot
"""

params_beamline = {
    'h5data': '/a/scan_data/data_04',
    'instrument': 'Cristal@Soleil',
    'monitor': None,
    'pixelsize': 55e-6,
    'ptychomotors': ['/a/scan_data/trajectory_1_1', '/a/scan_data/trajectory_1_2'],
    'saveprefix': '{data_prefix}-{scan:04d}/Run{run:04d}'
}

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


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

    def load_scan(self):
        h5 = h5py.File(self.params['data'], 'r')

        xmot, ymot = self.params['ptychomotors'].split(',')[0:2]

        self.x, self.y = np.array(h5[xmot][()]), np.array(h5[ymot][()])

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

        if self.params['monitor'] is not None:
            mon = np.array(h5[self.params['monitor']][()])
            self.raw_data_monitor = mon.copy()
            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*<I_obs>'
                    ' 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)

        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]
        h5.close()
        self.imgn = imgn

    def load_data(self):
        h5 = h5py.File(self.params['data'], 'r')
        # Load all frames
        imgn = self.imgn
        t0 = timeit.default_timer()
        vimg = None
        d0 = 0

        sys.stdout.write('Reading HDF5 frames: ')
        # frames are grouped in different subentries
        i0 = 0
        entry0 = 1
        h5entry = self.params["h5data"]
        print("\nReading h5 data entry: %s" % (h5entry))
        h5d = np.array(h5[h5entry][()])
        ii = 0
        for i in imgn:
            if (i - imgn[0]) % 20 == 0:
                sys.stdout.write('%d ' % (i - imgn[0]))
                sys.stdout.flush()
            # Load all frames
            if i >= (i0 + len(h5d)):
                # Read next data pack
                i0 += len(h5d)
                entry0 += 1
                h5entry = self.params["h5data"] % entry0
                print("\nReading h5 data entry: %s" % (h5entry))
                h5d = np.array(h5[h5entry][()])
            frame = h5d[i - i0]
            if vimg is None:
                vimg = np.empty((len(imgn), frame.shape[0], frame.shape[1]), dtype=frame.dtype)
            vimg[ii] = frame
            d0 += frame
            ii += 1
        print("\n")
        dt = timeit.default_timer() - t0
        print('Time to read all frames: %4.1fs [%5.2f Mpixel/s]' % (dt, d0.size * len(vimg) / 1e6 / dt))
        if self.raw_mask is not None:
            if self.raw_mask.sum() > 0:
                print("\nMASKING %d pixels from detector flags" % (self.raw_mask.sum()))

        self.raw_data = vimg
        self.load_data_post_process()


class PtychoRunnerCristal(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 = PtychoRunnerScanCristal

    @classmethod
    def make_parser(cls, default_par,  description=None, script_name="pynx-ptycho-simulation", 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)")
        p = default_par

        parser = super().make_parser(default_par, script_name, description, epilog)
        grp = parser.add_argument_group("CXI parameters")
        grp.add_argument('--data', type=str, default=None, required=True,
                         help='path to the nexus scan data file, e.g. /some/dir/to/data.h5')

        grp.add_argument('--h5data', '--h5path', type=str, default=p['h5data'],
                         required=True, help='entry/data/data_%%6d: generic hdf5 path '
                                             'to the stack of 2d images inside hdf5 file')

        grp.add_argument('--ptycho_motors', '--ptychomotors', type=str, default=None, required=True,
                         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..). Values will be read "
                              "from the hdf5 files, and are assumed to be in microns.\n"
                              "Examples:\n"
                              "  * --ptychomotors pix piz\n"
                              "  * --ptychomotors pix piz -x y\n"
                              "Note that if the '--xy -y x' command-line argument is used, "
                              "it is applied _after_ this, using '--ptychomotors pix piz -x y' "
                              "is equivalent to '--ptychomotors pix piz --xy -x y'")

        grp.add_argument('--monitor', type=str, default=None,
                         help='hdf5 path to the monitor. If supplied, 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 PtychoRunnerCristal.make_parser(default_params)
