#!/usr/bin/env python
# -*- coding: utf-8 -*-
# ----------------------------------------------------------------------- #
# Copyright (c) 2025, UChicago Argonne, LLC. All rights reserved.         #
#                                                                         #
# Copyright 2025. UChicago Argonne, LLC. This software was produced       #
# under U.S. Government contract DE-AC02-06CH11357 for Argonne National   #
# Laboratory (ANL), which is operated by UChicago Argonne, LLC for the    #
# U.S. Department of Energy. The U.S. Government has rights to use,       #
# reproduce, and distribute this software.  NEITHER THE GOVERNMENT NOR    #
# UChicago Argonne, LLC MAKES ANY WARRANTY, EXPRESS OR IMPLIED, OR        #
# ASSUMES ANY LIABILITY FOR THE USE OF THIS SOFTWARE.  If software is     #
# modified to produce derivative works, such modified software should     #
# be clearly marked, so as not to confuse it with the version available   #
# from ANL.                                                               #
#                                                                         #
# Additionally, redistribution and use in source and binary forms, with   #
# or without modification, are permitted provided that the following      #
# conditions are met:                                                     #
#                                                                         #
#     * Redistributions of source code must retain the above copyright    #
#       notice, this list of conditions and the following disclaimer.     #
#                                                                         #
#     * Redistributions in binary form must reproduce the above copyright #
#       notice, this list of conditions and the following disclaimer in   #
#       the documentation and/or other materials provided with the        #
#       distribution.                                                     #
#                                                                         #
#     * Neither the name of UChicago Argonne, LLC, Argonne National       #
#       Laboratory, ANL, the U.S. Government, nor the names of its        #
#       contributors may be used to endorse or promote products derived   #
#       from this software without specific prior written permission.     #
#                                                                         #
# THIS SOFTWARE IS PROVIDED BY UChicago Argonne, LLC AND CONTRIBUTORS     #
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT       #
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS       #
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL UChicago     #
# Argonne, LLC OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,        #
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,    #
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;        #
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER        #
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT      #
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN       #
# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE         #
# POSSIBILITY OF SUCH DAMAGE.                                             #
# ----------------------------------------------------------------------- #
import abc
import copy

import numpy as np
from aps.common.fit.gaussian import calculate_1D_gaussian_fit
from aps.common.fit.exponentially_modified_gaussian import fit_emg_hybrid
from aps.common.fit.exponentially_modified_pseudo_voigt import fit_epv_hybrid

from aps.beamline_driver.beam_management.facade import BeamProperties, Beam, MethodCallingType
from aps.beamline_driver.beam_management.abstract_beam_manager import DirectBeamImageBeamManager, DirectBeamImageProperties

from aps.beamline_driver.common.utility import generate_noise, calculate_projections_over_noise

from aps.common.plot.image import get_average, get_fwhm, get_sigma, get_peak_location

from scipy.ndimage import gaussian_filter, uniform_filter
from hybrid_methods.coherence.hybrid_screen import StdIOHybridListener

try:    import bluesky.plan_stubs as bps
except: print("bluesky not installed")

NOISE_DEFAULT_VALUE = 100

class NullHybridListener(StdIOHybridListener):
    def status_message(self, message : str): pass
    def set_progress_value(self, percentage : int): pass


def convolve_wire(profile, coordinates, wire_width):
    dx = coordinates[1] - coordinates[0]

    nbox = int(round(wire_width / dx))
    if nbox % 2 == 0: nbox += 1

    box = np.ones(nbox) / nbox
    pad = nbox // 2

    profile_pad  = np.pad(profile, pad_width=pad, mode='reflect')
    profile_conv = np.convolve(profile_pad, box, mode='valid')

    return profile_conv

class S4BeamManager(DirectBeamImageBeamManager):
    def __init__(self, method_calling_type=MethodCallingType.STANDARD, **kwargs):
        super(S4BeamManager, self).__init__(method_calling_type=method_calling_type, **kwargs)

        arguments = copy.deepcopy(kwargs)

        self._to_micron              = arguments.get("to_micron", True)

        # IMAGE
        self._nbins_h                = arguments.get("nbins_h", 201)
        self._nbins_v                = arguments.get("nbins_v", 201)
        self._xrange                 = arguments.get("xrange", None)
        self._yrange                 = arguments.get("yrange", None)
        self._blur_image             = arguments.get("blur_image", False)

        # SMART SCAN
        self._scan_mode              = arguments.get("scan_mode", False)
        self._scan_h                 = arguments.get("scan_h", [-10.0, 10.0, 401])
        self._scan_v                 = arguments.get("scan_v", [-10.0, 10.0, 401])
        self._relative_scan_h        = arguments.get("relative_scan_h", [-1.0, 1.0, 201])
        self._relative_scan_v        = arguments.get("relative_scan_v", [-1.0, 1.0, 201])
        self._q_h                    = arguments.get("q_h", 0.0)
        self._q_v                    = arguments.get("q_v", 0.0)
        self._motor_shift_signs      = arguments.get("motor_shift_signs",  {})
        self._activate_smart_scan    = arguments.get("activate_smart_scan", False)
        self._use_fit                = arguments.get("use_fit_as_beam_properties", 0)
        self._wire_width             = arguments.get("wire_width", None)

        # NOISE
        self._add_noise              = arguments.get("add_noise", False)
        self._noise                  = arguments.get("noise", NOISE_DEFAULT_VALUE if self._add_noise else None)
        self._percentage_fluctuation = arguments.get("percentage_fluctuation", 10.0) * 1e-2
        self._calculate_over_noise   = arguments.get("calculate_over_noise", False)
        self._noise_threshold        = arguments.get("noise_threshold", 1.5)

        if self._to_micron:
            if not self._xrange is None: self._xrange = (np.array(self._xrange) * 1e-6).tolist()
            if not self._yrange is None: self._yrange = (np.array(self._yrange) * 1e-6).tolist()

            if not self._scan_h is None:
                self._scan_h[0] *= 1e-6
                self._scan_h[1] *= 1e-6
            if not self._scan_v is None:
                self._scan_v[0] *= 1e-6
                self._scan_v[1] *= 1e-6
            if not self._relative_scan_h is None:
                self._relative_scan_h[0] *= 1e-6
                self._relative_scan_h[1] *= 1e-6
            if not self._relative_scan_v is None:
                self._relative_scan_v[0] *= 1e-6
                self._relative_scan_v[1] *= 1e-6
        else:
            if not self._wire_width is None: self._wire_width *= 1e6

        self._last_centroid_h = None
        self._last_centroid_v = None

    def set_method_calling_type(self, value):
        super(S4BeamManager, self).set_method_calling_type(value)

        if self.method_calling_type() == MethodCallingType.GENERATOR:
            setattr(self, "_set_initial_flux",  self.__set_initial_flux_generator)
            setattr(self, "_get_beam",          self.__get_beam_generator)
        else:
            setattr(self, "_set_initial_flux",  self.__set_initial_flux)
            setattr(self, "_get_beam",          self.__get_beam)

    def __set_initial_flux_generator(self, *args, **kwargs): yield from bps.null()

    def __set_initial_flux(self, *args, **kwargs): pass

    def __get_beam_generator(self, *args, **kwargs) -> Beam:
        yield from bps.null()

        self.__is_first_trial    = kwargs.get("is_first_trial", True)
        self.__delta_positions   = kwargs.get("delta_positions", {})
        self.__current_positions = kwargs.get("current_positions", {})

        return Beam(self._get_simulated_beam(**kwargs))

    def __get_beam(self, *args, **kwargs) -> Beam:
        self.__is_first_trial    = kwargs.get("is_first_trial", True)
        self.__delta_positions   = kwargs.get("delta_positions", {})
        self.__current_positions = kwargs.get("current_positions", {})

        return Beam(self._get_simulated_beam(**kwargs))

    def get_beam_properties(self, beam : Beam, **kwargs) -> BeamProperties:
        if not self._scan_mode:
            beam_properties = super(S4BeamManager, self).get_beam_properties(beam, **kwargs)
        else:
            beam_histogram  = self._get_beam_histogram(beam, **kwargs)
            beam_data       = beam_histogram.to_dictionary()
            beam_properties = BeamProperties()

            if self._wire_width is not None:
                try:
                    profile_h = convolve_wire(beam_data["profile_h"], beam_data["coordinate_h"], self._wire_width)
                    profile_v = convolve_wire(beam_data["profile_v"], beam_data["coordinate_v"], self._wire_width)
                except:
                    profile_h = beam_data["profile_h"]
                    profile_v = beam_data["profile_v"]

                beam_data["profile_h"] = profile_h
                beam_data["profile_v"] = profile_v

            centroid_h           = get_average(beam_data["profile_h"], beam_data["coordinate_h"])
            peak_distance_h      = get_peak_location(beam_data["profile_h"], beam_data["coordinate_h"])
            fwhm_h, _, _         = get_fwhm(beam_data["profile_h"], beam_data["coordinate_h"])
            sigma_h              = get_sigma(beam_data["profile_h"], beam_data["coordinate_h"])
            peak_intensity_h     = np.max(beam_data["profile_h"])
            integral_intensity_h = np.sum(beam_data["profile_h"])

            beam_properties.set_parameter(DirectBeamImageProperties.CENTROID_H, centroid_h)
            beam_properties.set_parameter(DirectBeamImageProperties.PEAK_LOCATION_H, peak_distance_h)
            beam_properties.set_parameter(DirectBeamImageProperties.FWHM_H, fwhm_h)
            beam_properties.set_parameter(DirectBeamImageProperties.SIGMA_H, sigma_h)
            beam_properties.set_parameter("coordinate_h", beam_data["coordinate_h"])
            beam_properties.set_parameter(self.get_beam_histogram_name() + "_h", beam_data["profile_h"])

            centroid_v           = get_average(beam_data["profile_v"], beam_data["coordinate_v"])
            peak_distance_v      = get_peak_location(beam_data["profile_v"], beam_data["coordinate_v"])
            fwhm_v, _, _         = get_fwhm(beam_data["profile_v"], beam_data["coordinate_v"])
            sigma_v              = get_sigma(beam_data["profile_v"], beam_data["coordinate_v"])
            peak_intensity_v     = np.max(beam_data["profile_v"])
            integral_intensity_v = np.sum(beam_data["profile_v"])

            beam_properties.set_parameter(DirectBeamImageProperties.CENTROID_V, centroid_v)
            beam_properties.set_parameter(DirectBeamImageProperties.PEAK_LOCATION_V, peak_distance_v)
            beam_properties.set_parameter(DirectBeamImageProperties.FWHM_V, fwhm_v)
            beam_properties.set_parameter(DirectBeamImageProperties.SIGMA_V, sigma_v)
            beam_properties.set_parameter("coordinate_v", beam_data["coordinate_v"])
            beam_properties.set_parameter(self.get_beam_histogram_name() + "_v", beam_data["profile_v"])

            peak_intensity     = max(peak_intensity_h, peak_intensity_v)
            integral_intensity = max(integral_intensity_h, integral_intensity_v)

            beam_properties.set_parameter(DirectBeamImageProperties.PEAK_INTENSITY, peak_intensity)
            beam_properties.set_parameter(DirectBeamImageProperties.INTEGRAL_INTENSITY, integral_intensity)

            beam_properties.set_parameter("is_empty_beam", integral_intensity <= self._get_input_parameters().get_parameter("empty_beam_threshold", 0.0))

            if self._use_fit > 0:
                if self._use_fit == 1:
                    try:    result_h = calculate_1D_gaussian_fit(beam_data["profile_h"], beam_data["coordinate_h"])
                    except: result_h = {"fwhm_x" : fwhm_h, "sigma_x" : sigma_h, "center_x" : peak_distance_h, "amplitude" : peak_intensity_h, "offset" : self._noise if self._add_noise else 0}
                    try:    result_v = calculate_1D_gaussian_fit(beam_data["profile_v"], beam_data["coordinate_v"])
                    except: result_v = {"fwhm_x" : fwhm_v, "sigma_x" : sigma_v, "center_x" : peak_distance_v, "amplitude" : peak_intensity_v, "offset" : self._noise if self._add_noise else 0}

                    beam_properties.set_parameter("gaussian_fit_h", result_h)
                    beam_properties.set_parameter("gaussian_fit_v", result_v)

                    fwhm_h = result_h["fwhm_x"]
                    fwhm_v = result_v["fwhm_x"]
                    sigma_h = result_h["sigma_x"]
                    sigma_v = result_v["sigma_x"]
                    center_h = result_h["center_x"]
                    center_v = result_v["center_x"]
                    peak_intensity = 0.5 * result_h["amplitude"] + 0.5 * result_v["amplitude"]

                    print("from gaussian fit: center (h x v):", center_h, center_v)
                elif self._use_fit == 2:
                    try:    result_h = fit_emg_hybrid(beam_data["coordinate_h"], beam_data["profile_h"])
                    except: result_h = {"popt" : [peak_intensity_h, peak_distance_h, sigma_h, 0.0, 0.0],
                                        "yhat": np.zeros(beam_data["profile_h"].size)}
                    try:    result_v = fit_emg_hybrid(beam_data["coordinate_v"], beam_data["profile_v"])
                    except: result_v = {"popt" : [peak_intensity_v, peak_distance_v, sigma_v, 0.0, 0.0],
                                        "yhat": np.zeros(beam_data["profile_v"].size)}

                    beam_properties.set_parameter("emg_fit_h", result_h)
                    beam_properties.set_parameter("emg_fit_v", result_v)

                    A, mu, sigma, tau, y0 = result_h["p_opt"]

                    fwhm_h = 2.355 * sigma
                    sigma_h = sigma
                    center_h = mu
                    peak_intensity = 0.5 * A

                    A, mu, sigma, tau, y0 = result_v["p_opt"]

                    fwhm_v = 2.355 * sigma
                    sigma_v = sigma
                    center_v = mu
                    peak_intensity += 0.5 * A

                    print("from emg fit: center (h x v):", center_h, center_v)
                elif self._use_fit == 3:
                    try:    result_h = fit_epv_hybrid(beam_data["coordinate_h"], beam_data["profile_h"])
                    except: result_h = {"popt" : [peak_intensity_h, peak_distance_h, sigma_h*2.355, 1.0, 0.0, 0.0],
                                        "yhat": np.zeros(beam_data["profile_h"].size)}
                    try:    result_v = fit_epv_hybrid(beam_data["coordinate_v"], beam_data["profile_v"])
                    except: result_v = {"popt" : [peak_intensity_v, peak_distance_v, sigma_v*2.355, 1.0, 0.0, 0.0],
                                        "yhat": np.zeros(beam_data["profile_v"].size)}

                    beam_properties.set_parameter("epv_fit_h", result_h)
                    beam_properties.set_parameter("epv_fit_v", result_v)

                    A, mu, fwhm, eta, tau, y0 = result_h["p_opt"]

                    fwhm_h = fwhm
                    sigma_h = fwhm / 2.355
                    center_h = mu
                    peak_intensity = 0.5 * A

                    A, mu, fwhm, eta, tau, y0 = result_v["p_opt"]

                    fwhm_v = fwhm
                    sigma_v = fwhm / 2.355
                    center_v = mu
                    peak_intensity += 0.5 * A

                    print("from epv fit: center (h x v):", center_h, center_v)

                beam_properties.set_parameter(DirectBeamImageProperties.PEAK_LOCATION_H, center_h)
                beam_properties.set_parameter(DirectBeamImageProperties.PEAK_LOCATION_V, center_v)
                beam_properties.set_parameter(DirectBeamImageProperties.FWHM_H, fwhm_h)
                beam_properties.set_parameter(DirectBeamImageProperties.FWHM_V, fwhm_v)
                beam_properties.set_parameter(DirectBeamImageProperties.SIGMA_H, sigma_h)
                beam_properties.set_parameter(DirectBeamImageProperties.SIGMA_V, sigma_v)
                beam_properties.set_parameter(DirectBeamImageProperties.PEAK_INTENSITY, peak_intensity)

            self._last_centroid_h = beam_properties.get_parameter(DirectBeamImageProperties.CENTROID_H) * (1e-6 if self._to_micron else 1.0)
            self._last_centroid_v = beam_properties.get_parameter(DirectBeamImageProperties.CENTROID_V) * (1e-6 if self._to_micron else 1.0)

            self._plot_beam(beam_properties, **kwargs)

        return beam_properties

    @abc.abstractmethod
    def _get_simulated_beam(self) -> Beam: raise NotImplementedError

    def _get_beam_histogram(self, beam : Beam, **kwargs) -> BeamProperties:
        s4_beam = beam.data

        if not self._scan_mode:
            ticket = s4_beam.histo2(col_h=1,
                                    col_v=3,
                                    nbins_h=int(self._nbins_h),
                                    nbins_v=int(self._nbins_v),
                                    nolost=1,
                                    xrange=self._xrange,
                                    yrange=self._yrange)

            if not kwargs.get("blur_image", False):  histogram = ticket["histogram"]
            else:
                filter = kwargs.get("filter", "gaussian")

                if   filter == "gaussian": histogram = gaussian_filter(ticket["histogram"], sigma=kwargs.get("sigma", 2))
                elif filter == "uniform":  histogram = uniform_filter(ticket["histogram"],  size=kwargs.get("sigma", 5))
        else:
            if not self._activate_smart_scan or self.__is_first_trial:
                ticket = s4_beam.histo2(col_h=1,
                                        col_v=3,
                                        nbins_h=int(self._scan_h[2]),
                                        nbins_v=int(self._scan_v[2]),
                                        nolost=1,
                                        xrange=[self._scan_h[0], self._scan_h[1]],
                                        yrange=[self._scan_v[0], self._scan_v[1]])
                self._last_centroid_h = None
                self._last_centroid_v = None
            else:
                try:
                    center_h = (self._last_centroid_h +
                                self._get_dx_total_shift(self.__current_positions, self.__delta_positions))
                    if not np.isnan(center_h):
                        xrange  = [center_h + self._relative_scan_h[0], center_h + self._relative_scan_h[1]]
                        nbins_h = int(self._relative_scan_h[2])
                    else:
                        xrange  = [self._scan_h[0],  self._scan_h[1]]
                        nbins_h = int(self._scan_h[2])

                    center_v = (self._last_centroid_v +
                                self._get_dy_total_shift(self.__current_positions, self.__delta_positions))
                    if not np.isnan(center_v):
                        yrange  = [center_v + self._relative_scan_v[0], center_v + self._relative_scan_v[1]]
                        nbins_v = int(self._relative_scan_v[2])
                    else:
                        yrange  = [self._scan_v[0], self._scan_v[1]]
                        nbins_v = int(self._scan_v[2])
                except:
                    xrange  = [self._scan_h[0], self._scan_h[1]]
                    yrange  = [self._scan_v[0], self._scan_v[1]]
                    nbins_h = int(self._scan_h[2])
                    nbins_v = int(self._scan_v[2])

                ticket = s4_beam.histo2(col_h=1,
                                        col_v=3,
                                        nbins_h=nbins_h,
                                        nbins_v=nbins_v,
                                        nolost=1,
                                        xrange=xrange,
                                        yrange=yrange)

            histogram = ticket["histogram"]

        coordinate_h = ticket['bin_h_center'] * (1e6 if self._to_micron else 1.0)
        coordinate_v = ticket['bin_v_center'] * (1e6 if self._to_micron else 1.0)

        if self._add_noise:
            histogram = generate_noise(histogram, self._noise, self._percentage_fluctuation)

            if self._calculate_over_noise:
                _, histogram_h, histogram_v = calculate_projections_over_noise(histogram, self._noise_threshold)
            else:
                histogram_h = histogram.sum(axis=1)
                histogram_v = histogram.sum(axis=0)
        else:
            histogram_h = ticket['histogram_h']
            histogram_v = ticket['histogram_v']

        if self._scan_mode:
            return BeamProperties(coordinate_h=coordinate_h,
                                  coordinate_v=coordinate_v,
                                  profile_h=histogram_h,
                                  profile_v=histogram_v)
        else:
            return BeamProperties(histogram=histogram,
                                  coordinate_h=coordinate_h,
                                  coordinate_v=coordinate_v,
                                  histogram_h=histogram_h,
                                  histogram_v=histogram_v)

    def _plot_beam(self, beam_properties, **kwargs):
        kwargs["use_fit"] = self._use_fit
        super(S4BeamManager, self)._plot_beam(beam_properties, **kwargs)

    @abc.abstractmethod
    def _get_dx_total_shift(self, current_positions, delta_positions): raise NotImplementedError

    @abc.abstractmethod
    def _get_dy_total_shift(self, current_positions, delta_positions): raise NotImplementedError

