import os, sys
import numpy

from AnyQt.QtWidgets import QLabel, QMessageBox, QSizePolicy
from AnyQt.QtGui import QTextCursor, QPixmap
from AnyQt.QtCore import Qt

from shadow4.physical_models.prerefl.prerefl import PreRefl

import orangecanvas.resources as resources

from orangewidget import gui
from orangewidget.settings import Setting
from orangewidget.widget import Output

from oasys2.widget.widget import OWWidget, OWAction
from oasys2.widget import gui as oasysgui
from oasys2.widget.util import congruence
from oasys2.canvas.util.canvas_util import add_widget_parameters_to_module
from oasys2.widget.util.widget_util import EmittingStream

from dabax.dabax_xraylib import DabaxXraylib

from orangecontrib.shadow4.util.shadow4_objects import PreReflPreProcessorData
from orangecontrib.shadow4.util.shadow4_util import ShadowPhysics
from orangecontrib.shadow4.widgets.gui.plots import plot_data1D, plot_data2D, plot_multi_data1D
from orangecontrib.shadow4.util.python_script import PythonScript

XRAYLIB_AVAILABLE = True

try: import xraylib
except: XRAYLIB_AVAILABLE = False

from dabax.dabax_files import dabax_f1f2_files, dabax_crosssec_files

class OWPrerefl(OWWidget):
    name = "PreRefl (absorbers, mirrors and lenses)"
    id = "PreRefl"
    description = "Calculation of mirror reflectivity profile"
    icon = "icons/prerefl.png"
    author = "create_widget.py"
    maintainer_email = "srio@esrf.eu"
    priority = 10
    category = ""
    keywords = ["xoppy", "xsh_prerefl"]

    class Outputs:
        preprocessor_data = Output("PreRefl PreProcessor Data", PreReflPreProcessorData, default=True, auto_summary=False)

    want_main_area = True

    symbol = Setting("SiC")
    density = Setting(3.217)
    prerefl_file = Setting("reflec.dat")
    e_min = Setting(100.0)
    e_max = Setting(20000.0)
    e_step = Setting(100.0)

    #
    # Plots
    #
    plot_flag = Setting(0)
    scan_e_n = Setting(100)
    scan_e_from = Setting(5000.0)
    scan_e_to = Setting(10000.0)
    scan_a0 = Setting(5.0)

    scan_a_n = Setting(100)
    scan_a_from = Setting(0)
    scan_a_to = Setting(10.0)
    scan_e0 = Setting(8000.0)

    MATERIAL_CONSTANT_LIBRARY_FLAG = Setting(1)
    DABAX_F1F2_FILE_INDEX = Setting(0)
    DABAX_CROSSSEC_FILE_INDEX = Setting(0)

    IMAGE_WIDTH  = 860
    IMAGE_HEIGHT = 545

    MAX_WIDTH          = 1320
    MAX_HEIGHT         = 720
    CONTROL_AREA_WIDTH = 405
    TABS_AREA_HEIGHT   = 615

    usage_path = os.path.join(resources.package_dirname("orangecontrib.shadow4.widgets.gui") , "misc", "prerefl_usage.png")

    prerefl_instance = None

    def __init__(self):
        super().__init__()

        self.runaction = OWAction("Compute", self)
        self.runaction.triggered.connect(self.compute)
        self.addAction(self.runaction)

        self.setFixedWidth(self.MAX_WIDTH)
        self.setFixedHeight(self.MAX_HEIGHT)

        gui.separator(self.controlArea)

        box0 = gui.widgetBox(self.controlArea, "",orientation="horizontal")

        button = gui.button(box0, self, "Compute", callback=self.compute)
        button.setFixedHeight(35)

        gui.separator(self.controlArea)

        tabs_setting = oasysgui.tabWidget(self.controlArea)
        tabs_setting.setFixedHeight(self.TABS_AREA_HEIGHT)
        tabs_setting.setFixedWidth(self.CONTROL_AREA_WIDTH-5)

        self.main_tabs = oasysgui.tabWidget(self.mainArea)
        tab_out = oasysgui.createTabPage(self.main_tabs, "Output")
        self.plot_tab = oasysgui.createTabPage(self.main_tabs, "Plots (optional scan)")

        ##############
        # script tab
        script_tab = oasysgui.createTabPage(self.main_tabs, "Script")
        self.shadow4_script = PythonScript()
        self.shadow4_script.code_area.setFixedHeight(400)

        script_box = gui.widgetBox(script_tab, "Python script", addSpace=True, orientation="horizontal")
        script_box.layout().addWidget(self.shadow4_script)

        ##############

        tab_bas = oasysgui.createTabPage(tabs_setting, "Settings")

        tab_dabax = oasysgui.createTabPage(tabs_setting, "Advanced")

        tab_input_plots = oasysgui.createTabPage(tabs_setting, "Plots")

        tab_usa = oasysgui.createTabPage(tabs_setting, "Use of the Widget")

        self.populate_tab_basic_settings(tab_bas)

        self.populate_tab_dabax(tab_dabax)

        self.populate_tab_plots(tab_input_plots)

        self.populate_tab_use_of_widget(tab_usa)

        self.shadow_output = oasysgui.textArea()
        out_box = oasysgui.widgetBox(tab_out, "System Output", addSpace=True, orientation="horizontal", height=self.TABS_AREA_HEIGHT)
        out_box.layout().addWidget(self.shadow_output)

        self.process_showers()

        gui.rubber(self.controlArea)

        self.set_visibility()

    def populate_tab_basic_settings(self, tab_bas):
        box = oasysgui.widgetBox(tab_bas, "Parameters", orientation="vertical")
        idx = -1

        # widget index 0
        idx += 1
        oasysgui.lineEdit(box, self, "symbol", tooltip="symbol",
                          label=self.unitLabels()[idx], addSpace=True, labelWidth=200, orientation="horizontal",
                          callback=self.set_Density)
        self.show_at(self.unitFlags()[idx], box)

        # widget index 1
        idx += 1
        oasysgui.lineEdit(box, self, "density", tooltip="density",
                          label=self.unitLabels()[idx], addSpace=True, valueType=float, labelWidth=200,
                          orientation="horizontal")
        self.show_at(self.unitFlags()[idx], box)

        # widget index 2
        idx += 1
        box_2 = oasysgui.widgetBox(box, "", addSpace=True, orientation="horizontal")

        self.le_prerefl_file = oasysgui.lineEdit(box_2, self, "prerefl_file", tooltip="prerefl_file",
                                                 label=self.unitLabels()[idx], addSpace=True, labelWidth=180,
                                                 orientation="horizontal")

        gui.button(box_2, self, "...", callback=self.selectFile)

        self.show_at(self.unitFlags()[idx], box)

        # widget index 3
        idx += 1
        oasysgui.lineEdit(box, self, "e_min", tooltip="e_min",
                          label=self.unitLabels()[idx], addSpace=True,
                          valueType=float, labelWidth=200, orientation="horizontal")
        self.show_at(self.unitFlags()[idx], box)

        # widget index 4
        idx += 1
        oasysgui.lineEdit(box, self, "e_max", tooltip="e_max",
                          label=self.unitLabels()[idx], addSpace=True,
                          valueType=float, labelWidth=200, orientation="horizontal")
        self.show_at(self.unitFlags()[idx], box)

        # widget index 5
        idx += 1
        oasysgui.lineEdit(box, self, "e_step", tooltip="e_step",
                          label=self.unitLabels()[idx], addSpace=True,
                          valueType=float, labelWidth=200, orientation="horizontal")
        self.show_at(self.unitFlags()[idx], box)

    def populate_tab_dabax(self, tab_dabax):
        # xraylib/dabax
        # widget index xx
        box1 = gui.widgetBox(tab_dabax, "Material Library")
        self.cb_material_library = \
            gui.comboBox(box1, self, "MATERIAL_CONSTANT_LIBRARY_FLAG",
                         label='Material Library', addSpace=True,
                         items=["xraylib", "dabax [default]"],
                         orientation="horizontal",
                         callback=self.set_visibility)

        # widget index xx
        self.dabax_f1f2_box = gui.widgetBox(box1)
        gui.comboBox(self.dabax_f1f2_box, self, "DABAX_F1F2_FILE_INDEX",
                     label="dabax f1f2 file", addSpace=True,
                     items=dabax_f1f2_files(),
                     orientation="horizontal")

        # widget index xx
        self.dabax_crosssec_box = gui.widgetBox(box1)
        gui.comboBox(self.dabax_crosssec_box, self, "DABAX_CROSSSEC_FILE_INDEX",
                     label="dabax cross sec file", addSpace=True,
                     items=dabax_crosssec_files(),
                     orientation="horizontal")

        if not XRAYLIB_AVAILABLE:
            self.MATERIAL_CONSTANT_LIBRARY_FLAG = 1
            self.cb_material_library.setEnabled(False)
        else:
            self.cb_material_library.setEnabled(True)

    def populate_tab_plots(self, tab_plots):

        box = gui.widgetBox(tab_plots, "Optional scan plots", orientation="vertical")

        gui.comboBox(box, self, "plot_flag", tooltip="plot_flag",
                     label="Scan plot", addSpace=True,
                     items=['No', 'refraction index', 'attenuation coefficient',
                            'mirror reflectivity energy-scan',
                            'mirror reflectivity grazing angle-scan',
                            'mirror reflectivity energy-angle-scan'],
                     valueType=int, orientation="horizontal", labelWidth=270, callback=self.do_plots)

        self.box_plots = gui.widgetBox(tab_plots, "Optional scan plots", orientation="vertical")

        self.box_plot_e = oasysgui.widgetBox(self.box_plots, "", addSpace=True, orientation="horizontal")
        oasysgui.lineEdit(self.box_plot_e, self, "scan_e_from", label="from [eV]", tooltip="scan_e_from",
                          valueType=float, addSpace=True, labelWidth=100, orientation="horizontal",
                          callback=self.do_plots)
        oasysgui.lineEdit(self.box_plot_e, self, "scan_e_to", label="to [eV]", tooltip="scan_e_to",
                          valueType=float, addSpace=True, labelWidth=20, orientation="horizontal",
                          callback=self.do_plots)
        oasysgui.lineEdit(self.box_plot_e, self, "scan_e_n", label="points", tooltip="scan_e_n",
                          valueType=int, addSpace=True, labelWidth=50, orientation="horizontal",
                          callback=self.do_plots)
        self.box_plot_a0 = oasysgui.widgetBox(self.box_plots, "", addSpace=True, orientation="horizontal")
        oasysgui.lineEdit(self.box_plot_a0, self, "scan_a0", label="Fixed angle [mdeg]", tooltip="scan_a0",
                          valueType=float, addSpace=True, labelWidth=200, orientation="horizontal",
                          callback=self.do_plots)

        self.box_plot_a = oasysgui.widgetBox(self.box_plots, "", addSpace=True, orientation="horizontal")
        oasysgui.lineEdit(self.box_plot_a, self, "scan_a_from", label="from [deg]", tooltip="scan_a_from",
                          valueType=float, addSpace=True, labelWidth=100, orientation="horizontal",
                          callback=self.do_plots)
        oasysgui.lineEdit(self.box_plot_a, self, "scan_a_to", label="to [deg]", tooltip="scan_a_to",
                          valueType=float, addSpace=True, labelWidth=20, orientation="horizontal",
                          callback=self.do_plots)
        oasysgui.lineEdit(self.box_plot_a, self, "scan_a_n", label="points", tooltip="scan_a_n",
                          valueType=int, addSpace=True, labelWidth=50, orientation="horizontal",
                          callback=self.do_plots)
        self.box_plot_e0 = oasysgui.widgetBox(self.box_plots, "", addSpace=True, orientation="horizontal")
        oasysgui.lineEdit(self.box_plot_e0, self, "scan_e0", label="Fixed energy [eV]", tooltip="scan_e0",
                          valueType=float, addSpace=True, labelWidth=200, orientation="horizontal",
                          callback=self.do_plots)


    def populate_tab_use_of_widget(self, tab_usa):
        tab_usa.setStyleSheet("background-color: white;")
        usage_box = oasysgui.widgetBox(tab_usa, "", addSpace=True, orientation="horizontal")
        label = QLabel("")
        label.setAlignment(Qt.AlignCenter)
        label.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
        label.setPixmap(QPixmap(self.usage_path))
        usage_box.layout().addWidget(label)

    def set_visibility(self):
        self.box_plots.setVisible(self.plot_flag > 0)
        self.box_plot_e.setVisible(False)
        self.box_plot_a0.setVisible(False)
        self.box_plot_a.setVisible(False)
        self.box_plot_e0.setVisible(False)
        self.dabax_f1f2_box.setVisible(False)
        self.dabax_crosssec_box.setVisible(False)


        if self.plot_flag == 1:
            self.box_plot_e.setVisible(True)
        elif self.plot_flag == 2:
            self.box_plot_e.setVisible(True)
        elif self.plot_flag == 3:
            self.box_plot_e.setVisible(True)
            self.box_plot_a0.setVisible(True)
        elif self.plot_flag == 4:
            self.box_plot_a.setVisible(True)
            self.box_plot_e0.setVisible(True)
        elif self.plot_flag == 5:
            self.box_plot_a.setVisible(True)
            self.box_plot_e.setVisible(True)

        if self.MATERIAL_CONSTANT_LIBRARY_FLAG > 0:
            self.dabax_f1f2_box.setVisible(True)
            self.dabax_crosssec_box.setVisible(True)

    def unitLabels(self):
         return ['Element/Compound formula','Density [ g/cm3 ]','File name (for SHADOW):','Minimum energy [eV]','Maximum energy [eV]','Energy step [eV]']

    def unitFlags(self):
         return ['True','True','True','True','True','True']

    def selectFile(self):
        self.le_prerefl_file.setText(oasysgui.selectFileFromDialog(self, self.prerefl_file, "Select Output File", file_extension_filter="Data Files (*.dat)"))

    def set_Density(self):
        if not self.symbol is None:
            if not self.symbol.strip() == "":
                self.symbol = self.symbol.strip()
                self.density = ShadowPhysics.getMaterialDensity(self.symbol)


    def compute(self):
        if self.MATERIAL_CONSTANT_LIBRARY_FLAG == 0:
            import xraylib
            material_constants_library = xraylib
            material_constants_library_str = "xraylib"
            material_constants_library_import = "\nimport xraylib"
        else:
            material_constants_library = DabaxXraylib(file_f1f2=dabax_f1f2_files()[self.DABAX_F1F2_FILE_INDEX],
                                                      file_CrossSec=dabax_crosssec_files()[self.DABAX_CROSSSEC_FILE_INDEX])
            material_constants_library_str = 'DabaxXraylib(file_f1f2="%s",file_CrossSec="%s")' % \
                                             (dabax_f1f2_files()[self.DABAX_F1F2_FILE_INDEX],
                                              dabax_crosssec_files()[self.DABAX_CROSSSEC_FILE_INDEX])
            material_constants_library_import = "\nfrom dabax.dabax_xraylib import DabaxXraylib"

        try:
            sys.stdout = EmittingStream(textWritten=self.writeStdOut)

            self.checkFields()

            PreRefl.prerefl(interactive=False,
                            SYMBOL=self.symbol,
                            DENSITY=self.density,
                            FILE=congruence.checkFileName(self.prerefl_file),
                            E_MIN=self.e_min,
                            E_MAX=self.e_max,
                            E_STEP=self.e_step,
                            materials_library=material_constants_library)

            self.Outputs.preprocessor_data.send(PreReflPreProcessorData(prerefl_data_file=self.prerefl_file))

            self.prerefl_instance = PreRefl()
            self.prerefl_instance.read_preprocessor_file(self.prerefl_file)
            self.prerefl_instance.preprocessor_info()

            script = "# script to create the PreRefl preprocessor file (for absorbers, mirrors and lenses)"
            script += material_constants_library_import
            script += "\nfrom shadow4.physical_models.prerefl.prerefl import PreRefl"
            script += '\nPreRefl.prerefl(interactive=False, SYMBOL="%s", DENSITY=%f, FILE="%s", E_MIN=%f, E_MAX=%f, E_STEP=%f, materials_library=%s)' % \
                   (self.symbol, self.density, self.prerefl_file, self.e_min, self.e_max, self.e_step, material_constants_library_str)
            script +="\n\n"

            script_plot = self.do_plots()

            self.shadow4_script.set_code(script + script_plot)

        except Exception as exception:
            QMessageBox.critical(self, "Error", str(exception), QMessageBox.Ok)
            if self.IS_DEVELOP: raise exception

    def checkFields(self):
        self.symbol = ShadowPhysics.checkCompoundName(self.symbol)
        self.density = congruence.checkStrictlyPositiveNumber(self.density, "Density")
        self.e_min  = congruence.checkPositiveNumber(self.e_min , "Minimum Energy")
        self.e_max  = congruence.checkStrictlyPositiveNumber(self.e_max , "Maximum Energy")
        self.e_step = congruence.checkStrictlyPositiveNumber(self.e_step, "Energy step")
        congruence.checkLessOrEqualThan(self.e_min, self.e_max, "Minimum Energy", "Maximum Energy")
        congruence.checkDir(self.prerefl_file)

    def do_plots(self):

        script = ""
        if self.plot_flag != 0:
            script += "\n# test plot\nfrom srxraylib.plot.gol import plot, plot_image"
            script += "\nimport numpy as np"
            script += "\n\nprerefl_instance = PreRefl()"
            script += "\nprerefl_instance.read_preprocessor_file('%s')" % self.prerefl_file
            script += "\n"

        self.set_visibility()

        self.plot_tab.layout().removeItem(self.plot_tab.layout().itemAt(0))

        if self.prerefl_instance is None: return

        if self.plot_flag == 0:
            plot_widget_id = plot_data1D([0], [0], xtitle="", ytitle="")
        elif self.plot_flag == 1:
            energy_array = numpy.linspace(self.scan_e_from, self.scan_e_to, self.scan_e_n)
            refraction_index = self.prerefl_instance.get_refraction_index(energy_array)
            delta = 1.0 - refraction_index.real
            beta = refraction_index.imag

            plot_widget_id = plot_multi_data1D(energy_array, [delta, beta],
                                         xtitle="Photon energy [eV]",
                                         ytitle="(n = 1 - delta + i beta)",
                                         ytitles=['delta', 'beta'])

            script += "\nenergy_array = np.linspace(%f, %f, %d)" % (self.scan_e_from, self.scan_e_to, self.scan_e_n)
            script += "\nrefraction_index = prerefl_instance.get_refraction_index(energy_array)"
            script += "\ndelta = 1.0 - refraction_index.real"
            script += "\nbeta = refraction_index.imag"
            script += "\nplot(energy_array, delta, energy_array, beta, legend=['delta','beta'], xtitle='Photon energy [eV]')"

        elif self.plot_flag == 2:
            energy_array = numpy.linspace(self.scan_e_from, self.scan_e_to, self.scan_e_n)
            att = self.prerefl_instance.get_attenuation_coefficient(energy_array)

            plot_widget_id = plot_data1D(energy_array, att,
                                         xtitle="Photon energy [eV]",
                                         ytitle="mu - attenuation coefficient cm^-1",
                                         )
            script += "\nenergy_array = np.linspace(%f, %f, %d)" % (self.scan_e_from, self.scan_e_to, self.scan_e_n)
            script += "\natt = prerefl_instance.get_attenuation_coefficient(energy_array)"
            script += "\nplot(energy_array, att, ytitle='mu - attenuation coefficient cm^-1', xtitle='Photon energy [eV]')"

        elif self.plot_flag == 3:
            energy_array = numpy.linspace(self.scan_e_from, self.scan_e_to, self.scan_e_n)
            RS, RP, _ = self.prerefl_instance.reflectivity_fresnel(grazing_angle_mrad=self.scan_a0,
                                                                           photon_energy_ev=energy_array,
                                                                           roughness_rms_A=0.0,
                                                                           )

            plot_widget_id = plot_multi_data1D(energy_array, [RS, RP],
                                         xtitle="Photon energy [eV]",
                                         ytitle="Mirror reflectivity @ %.3f mrad" % self.scan_a0,
                                         ytitles=['S-polarized', 'P-polarized'])

            script += "\nenergy_array = np.linspace(%f, %f, %d)" % (self.scan_e_from, self.scan_e_to, self.scan_e_n)
            script += "\nRS, RP, _ = prerefl_instance.reflectivity_fresnel(grazing_angle_mrad=%f, photon_energy_ev=energy_array, roughness_rms_A=0.0)" % \
                        (self.scan_a0)
            script += "\nplot(energy_array, RS, energy_array, RP, legend=['S-polarized', 'P-polarized'], xtitle='Photon energy [eV]', ytitle='Mirror reflectivity @ %.3f mrad')" % \
                      (self.scan_a0)

        elif self.plot_flag == 4:
            angle_array = numpy.linspace(self.scan_a_from, self.scan_a_to, self.scan_a_n)
            RS, RP, _ = self.prerefl_instance.reflectivity_fresnel(grazing_angle_mrad=angle_array,
                                                                           photon_energy_ev=self.scan_e0,
                                                                           roughness_rms_A=0.0,
                                                                           )

            plot_widget_id = plot_multi_data1D(angle_array, [RS, RP],
                                         xtitle="Grazing angle [mrad]",
                                         ytitle="Mirror reflectivity @ %.3f eV" % self.scan_e0,
                                         ytitles=['S-polarized', 'P-polarized'])

            script += "\nangle_array = np.linspace(%f, %f, %d)" % (self.scan_a_from, self.scan_a_to, self.scan_a_n)
            script += "\nRS, RP, _ = prerefl_instance.reflectivity_fresnel(grazing_angle_mrad=angle_array, photon_energy_ev=%f, roughness_rms_A=0.0)" % \
                        (self.scan_e0)
            script += "\nplot(angle_array, RS, angle_array, RP, legend=['S-polarized', 'P-polarized'], xtitle='Grazing angle [mrad]', ytitle='Mirror reflectivity @ %.3f eV')" % \
                      (self.scan_e0)

        elif self.plot_flag == 5:
            angle_array = numpy.linspace(self.scan_a_from, self.scan_a_to, self.scan_a_n)
            energy_array = numpy.linspace(self.scan_e_from, self.scan_e_to, self.scan_e_n)
            E = numpy.outer(energy_array, numpy.ones_like(angle_array))
            A = numpy.outer(numpy.ones_like(energy_array), angle_array)

            RS, RP, _ = self.prerefl_instance.reflectivity_fresnel(grazing_angle_mrad=A,
                                                                   photon_energy_ev=E,
                                                                   roughness_rms_A=0.0,
                                                                   )
            RS.shape = (energy_array.size, angle_array.size)

            plot_widget_id = plot_data2D(RS, energy_array, angle_array,
                                         xtitle="Photon energy [eV]",
                                         ytitle="Grazing angle [mrad]",
                                         title="Mirror reflectivity (S-pol)",
                                         )
            script += "\nangle_array = np.linspace(%f, %f, %d)" % (self.scan_a_from, self.scan_a_to, self.scan_a_n)
            script += "\nenergy_array = np.linspace(%f, %f, %d)" % (self.scan_e_from, self.scan_e_to, self.scan_e_n)
            script += "\nE = np.outer(energy_array, np.ones_like(angle_array))"
            script += "\nA = np.outer(np.ones_like(energy_array), angle_array)"
            script += "\nRS, RP, _ = prerefl_instance.reflectivity_fresnel(grazing_angle_mrad=A, photon_energy_ev=E, roughness_rms_A=0.0)"
            script += "\nplot_image(RS, energy_array, angle_array, xtitle='Photon energy [eV]', ytitle='Grazing angle [mrad]', title='Mirror reflectivity (S-pol)', aspect='auto')"

        self.plot_tab.layout().addWidget(plot_widget_id)
        return script

    def writeStdOut(self, text):
        cursor = self.shadow_output.textCursor()
        cursor.movePosition(QTextCursor.End)
        cursor.insertText(text)
        self.shadow_output.setTextCursor(cursor)
        self.shadow_output.ensureCursorVisible()

add_widget_parameters_to_module(__name__)

if __name__ == "__main__":
    import sys
    from AnyQt.QtWidgets import QApplication
    a = QApplication(sys.argv)
    ow = OWPrerefl()
    ow.show()
    a.exec()
    ow.saveSettings()