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

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

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

try:
    # Get the real number of processor cores available
    # os.sched_getaffinity is only available on some *nix platforms
    nproc = len(os.sched_getaffinity(0)) * psutil.cpu_count(logical=False) // psutil.cpu_count(logical=True)
except AttributeError:
    nproc = os.cpu_count()

helptext_epilog = """
###############################################################################
Examples:
    pynx-ptycho-p06 --data master.h5 --positions pos.h5 \
--probe focus,60e-6x200e-6,0.09 \
--algorithm analysis,ML**100,DM**200,nbprobe=3,probe=1 --verbose 10 --save all --saveplot --liveplot
"""

params_beamline = {
    'instrument': 'PO06@PetraIII'
}

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


def load_frame(filename, i, chunk_size, roi):
    import h5py as h5
    with h5.File(filename) as h:
        xmin, xmax, ymin, ymax = roi
        d = h[f'/entry/data/data_{i // chunk_size + 1:06d}'][i % chunk_size, ymin:ymax, xmin:xmax].astype(np.float32)
        # mask data above maximum allowed
        nmax = h['/entry/instrument/detector/bit_depth_image'][()]
        d[d > (2 ** nmax - 3)] = -1
    return d


def load_frame_kw(kw):
    return load_frame(**kw)


class PtychoRunnerScanP06(PtychoRunnerScan):

    def load_scan(self):
        if os.path.isfile(self.params['positions']) is False:
            raise PtychoRunnerException("positions hdf5 file does not exist: %s"
                                        % (self.params['positions']))
        with h5py.File(self.params['positions'], 'r') as h:
            x = h['/data/encoder_fast/data'][()] * 1e-6
            y = h['/data/encoder_slow/data'][()] * 1e-6

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

        imgn = np.arange(len(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)
            x = x.take(idx)
            y = 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]
                x = x[:N]
                y = y[:N]
        self.x, self.y, self.imgn = x, y, imgn

    def load_data(self):
        imgn = self.imgn
        if imgn is None:
            raise PtychoRunnerException("load_data(): imgn is None. Did you call load_scan() before ?")
        if os.path.isfile(self.params['data']) is False:
            raise PtychoRunnerException("data file does not exist: %s" % (self.params['data']))
        with h5py.File(self.params['data'], 'r') as h:
            if self.params['nrj'] is None:
                self.params['nrj'] = 12.3984 / h['/entry/instrument/beam/incident_wavelength'][()]
            if self.params['detectordistance'] is None:
                self.params['detectordistance'] = h['/entry/instrument/detector/detector_distance'][()]
                if self.params['detectordistance'] == 0:
                    raise PtychoRunnerException("/entry/instrument/detector/detector_distance is zero, "
                                                "please give the distance using --detector_distance")
            if self.params['pixelsize'] is None:
                self.params['pixelsize'] = h['/entry/instrument/detector/x_pixel_size'][()]

        # Frames are saved by chunks of n= 50 frames (typically or systematically ?)
        with h5py.File(self.params['data'], 'r') as h:
            chunk_size, ny, nx = h[f'/entry/data/data_{1:06d}'].shape
        if self.params['roi'] is None:
            warnings.warn('P06 data loading: no ROI was given - you may want to supply one if '
                          'many frames are loaded with a large detector')
            roi = 0, nx, 0, ny
        else:
            roi = self.params['roi']
            xmin, xmax, ymin, ymax = roi
            ny, nx = ymax - ymin, xmax - xmin
            self.roi_applied = True

        # Load all frames in //
        sys.stdout.write(f"Reading {len(imgn)} frames from HDF5 file [{nproc} // process]: 0")
        t0 = timeit.default_timer()
        sys.stdout.flush()
        vimg = np.empty((len(imgn), ny, nx), dtype=np.float32)
        with multiprocessing.Pool(nproc) as pool:
            vkw = [{'filename': self.params['data'], 'i': i, 'chunk_size': chunk_size,
                    'roi': roi} for i in self.imgn]
            results = pool.imap(load_frame_kw, vkw)  # , chunksize=1
            for i in range(len(vkw)):
                vimg[i] = results.next(timeout=20)
                if i % 20 == 0:
                    sys.stdout.write(f'.{i}')
                    sys.stdout.flush()
        print()
        dt = timeit.default_timer() - t0
        print('Time to read all frames: %4.1fs [%7.4f Go/s]' % (dt, vimg.nbytes / dt / 1024 ** 3))

        # if '/entry_1/instrument_1/detector_1/mask' in h:
        #     print("Loaded mask from data: /entry_1/instrument_1/detector_1/mask")
        #     self.raw_mask = h['/entry_1/instrument_1/detector_1/mask'][()]
        #     self.raw_mask = (self.raw_mask != 0).astype(np.int8)
        # if '/entry_1/instrument_1/detector_1/dark' in h:
        #     print("Loaded dark from data: /entry_1/instrument_1/detector_1/dark")
        #     self.dark = h['/entry_1/instrument_1/detector_1/dark'][()]
        self.raw_data = vimg
        self.load_data_post_process()


class PtychoRunnerP06(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 = PtychoRunnerScanP06

    @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 "
                           "at P06@PetraIII")

        parser = super().make_parser(default_par, script_name, description, epilog)
        grp = parser.add_argument_group("P06 parameters")
        grp.add_argument('--data', type=str, default=None, required=True,
                         help='master hdf5 file with all the detector data, '
                              'plus some metadata (pixel size, wavelength,..)')
        grp.add_argument('--positions', type=str, default=None, required=True,
                         help='hdf5 file with the motor positions')
        return parser


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