import os
import sys
import time

import matplotlib.pyplot as plt
import numpy as np

from .recipes import BeamWaist
from .rml import ObjectElement
from .simulate import Simulate


class PlotBeamwaist:
    """
    Trace and Export RawRaysOutgoing for optical elements, excluding image planes.
    """

    def __init__(self, directory: str, sim: Simulate):
        self.directory = "RAYPy_Simulation_" + directory
        self._original_directory = directory
        if isinstance(sim, Simulate):
            self._sim = sim
        self.lim = False
        self.step = False
        self.z = False
        self.step_z = False
        self.rot = False
        self.factor = False
        self.previous_results = False
        self.count_el = 0
        self.count_fig = 0

    def simulate_beamline(
        self, energy: float, /, source: ObjectElement = None, nrays: int = None, force: bool = False
    ):
        rp = BeamWaist(energy, source=source, nrays=nrays, sim_folder=self._original_directory)

        # test resolving power simulations
        self._sim.run(rp, multiprocessing=5, force=force)

    def _parse_beamline_elements(self, debug=False):
        self.element_names_list = []
        self.distance_list = []
        self.rotation_list = []
        for ind, oe in enumerate(self._sim.rml.beamline._children):
            for par in oe:
                try:
                    # append to distances
                    if ind != 0:
                        self.distance_list.append(
                            float(par.worldPosition.z.cdata) - sum(self.distance_list)
                        )
                    else:
                        self.distance_list.append(float(par.worldPosition.z.cdata))
                    # append to rotation
                    if ind == 0:
                        self.rotation_list.append(False)
                    elif hasattr(par, "azimuthalAngle"):
                        if par.azimuthalAngle.cdata == "0" or par.azimuthalAngle.cdata == "180":
                            self.rotation_list.append(False)
                        elif par.azimuthalAngle.cdata == "90" or par.azimuthalAngle.cdata == "270":
                            self.rotation_list.append(True)
                        else:
                            raise ValueError(
                                "Allowed azimuthal angles: 0,90,180,270:",
                                oe.name,
                                par.azimuthalAngle.cdata,
                            )
                    else:
                        self.rotation_list.append(False)
                    # append to element names list, only if it is not an imageplane
                    # if oe.get_attribute('type') != 'ImagePlaneBundle'
                    # and oe.get_attribute('type') != 'ImagePlane':
                    self.element_names_list.append(oe.attributes().original()["name"])
                except AttributeError:
                    print(oe.attributes().original()["name"], par.worldPosition.z.cdata)
                    pass
        if debug:
            print("DEBUG:: element_names_list", self.element_names_list)
            print("DEBUG:: distance_list", self.distance_list)
            print("DEBUG:: rotation_list", self.rotation_list)

    def trace_beamwaist(self, save_results: bool = True, element_names_list: list = None):
        self._parse_beamline_elements()
        # I overwrite the names for the moment, since I can not access them automatically
        if element_names_list is not None:
            self.element_names_list = element_names_list
        for ind in range(len(self.element_names_list)):
            if ind + 1 != len(self.element_names_list):
                self.add_element(
                    name=self.element_names_list[ind],
                    z=self.distance_list[ind + 1],
                    rot=self.rotation_list[ind],
                )
        if save_results is True:
            self.save_results()

    def reduce_Nrays(self, factor):
        self.factor = factor

    def load_previous_results(self, element_names_list=None):
        self._parse_beamline_elements()
        if element_names_list is not None:
            self.element_names_list = element_names_list
        start1 = time.time()
        print("Load results...")
        self.xh = np.loadtxt(os.path.join(self.directory, "xh.txt"))
        self.yh = np.loadtxt(os.path.join(self.directory, "yh.txt"))
        stop = time.time()
        print("time:", np.round(stop - start1, 2), "s")

    def save_results(self):
        np.savetxt(os.path.join(self.directory, "xh.txt"), self.xh)
        np.savetxt(os.path.join(self.directory, "yh.txt"), self.yh)

    def add_element(self, name, z, rot=False):
        start1 = time.time()
        z = np.arange(0, z + self.step_z, self.step_z)

        if self.previous_results is False:
            print("Tracing " + name)

            rays = np.loadtxt(
                os.path.join(self.directory, "round_0", "0_" + name + "-RawRaysOutgoing.csv"),
                skiprows=2,
            )
            if self.factor is not False:
                max_n_rays = int(rays.shape[0] / self.factor)
            if max_n_rays < 100:
                sys.exit("Set a lower reduction factor, there are no more rays to plot!")
            rays = rays[0:max_n_rays]

            if self.count_el == 0:
                self.xh, self.yh = self.trace(z, rays, rot)
                # print("DEBUG:: self.xh.shape, self.yh.shape:",self.xh.shape, self.yh.shape)
            else:
                txh, tyh = self.trace(z, rays, rot, name)
            if rot is True:
                txh = np.rot90(txh)
                tyh = np.rot90(tyh)
            if self.count_el != 0:
                # print("DEBUG:: txh.shape, tyh.shape:",txh.shape, tyh.shape)
                self.xh = np.concatenate((self.xh, txh), axis=1)
                self.yh = np.concatenate((self.yh, tyh), axis=1)

            stop = time.time()
            print("time:", np.round(stop - start1, 2), "s")
        self.count_el += 1

    def define_hist(self, lim, step):
        self.lim = lim
        self.step = step

    def define_zstep(self, step_z):
        self.step_z = step_z

    def trace_pos(self, rays, z):
        ox = 3
        oy = 4
        oz = 5
        dx = 6
        dy = 7
        dz = 8

        dxz = rays[:, dx] / rays[:, dz]
        x = rays[:, ox] + dxz * (z - rays[:, oz])
        dyz = rays[:, dy] / rays[:, dz]
        y = rays[:, oy] + dyz * (z - rays[:, oz])
        return x, y

    def make_histogram(self, x):
        xh = np.histogram(x, bins=np.arange(-self.lim - self.step, self.lim + self.step, self.step))
        return xh

    def trace(self, z_direction, rays, rot=False, element=None):
        for n, z in enumerate(z_direction):
            x, y = self.trace_pos(rays, z)
            shiftx = np.average(x)
            shifty = np.average(y)
            y = y - shifty
            x = x - shiftx
            if n == 0:
                xh = self.make_histogram(x)
                xh = xh[0]
                yh = self.make_histogram(y)
                yh = yh[0]
                np.argmax(yh)

            else:
                xh_temp = self.make_histogram(x)
                yh_temp = self.make_histogram(y)
                xh = np.concatenate((xh, xh_temp[0]), axis=0)
                yh = np.concatenate((yh, yh_temp[0]))
        xh = xh.reshape((n + 1, xh_temp[0].shape[0]))
        yh = yh.reshape((n + 1, yh_temp[0].shape[0]))

        xh = np.rot90(xh)
        if rot is False:
            return xh, np.rot90(yh)
        elif rot is True:
            return yh, np.rot90(np.flip(xh))  # no idea why I have to flip this

    def change_name(self, new_name, pos):
        self.element_names_list[pos] = new_name

    def plot(
        self,
        save_img=True,
        img_name="test",
        extension=".png",
        show_img=False,
        annotate_OE=False,
        lim_top=False,
        lim_side=False,
        debug=False,
    ):
        dx, dy = self.step_z, self.step
        xmax = self.yh.shape[1] * self.step_z + dx
        y, x = np.mgrid[slice(-self.lim, self.lim + dy, dy), slice(0, xmax, dx)]
        x = x[:, :-1]
        y = y[:, :-1]
        if debug:
            print("######################")
            print("Dimension check")
            print("yh shape:  ", self.yh.shape)
            print("xh shape:  ", self.xh.shape)
            print("x shape:    ", x.shape)
            print("y shape:    ", y.shape)
            print("x :    ", x)
            print(
                "self.yh.shape[1], self.step_z, dx",
                self.yh.shape,
                self.yh.shape[1],
                self.step_z,
                dx,
            )
            print("self.yh.shape[1]*self.step_z + dx", self.yh.shape[1] * self.step_z + dx)
            print("DEBUG:: self.distances", self.distances)
            print("DEBUG:: self.distance_list", self.distance_list)
            print("DEBUG:: self.elements", self.elements)
            print("DEBUG:: self.elements_name_list", self.element_names_list)

        g = 1.5
        plt.figure(self.count_fig)
        fig, (ax1, ax2) = plt.subplots(2, figsize=(6.4 * g, 4.8 * g))
        pcm = ax1.pcolormesh(x / 1000, y, self.xh, cmap="inferno")
        ax1.clear()
        ax1.pcolormesh(x / 1000, y, np.log(self.xh), cmap="inferno")
        ax2.pcolormesh(x / 1000, y, np.log(self.yh), cmap="inferno")

        ax1.set_title("top view")
        ax2.set_title("side view")
        ax2.set_xlabel("[meters]")
        ax1.set_ylabel("Beam Width [mm]")
        ax2.set_ylabel("Beam Width [mm]")
        cbar1 = fig.colorbar(pcm, ax=ax1)
        cbar2 = fig.colorbar(pcm, ax=ax2)
        cbar1.set_label("# of rays [a.u.]")  # , rotation=-90)
        cbar2.set_label("# of rays [a.u.]")
        posx = 0
        xtick_pos = []
        xtick_label = []
        # put correct ticks and labels
        plot_length = xmax / 1000
        beamline_length = sum(self.distance_list) / 1000
        beamline_el_pos = np.cumsum(self.distance_list) / 1000
        xtick_pos = (beamline_el_pos * plot_length) / beamline_length
        xtick_label = np.round(beamline_el_pos, 2)

        # label the elements
        if annotate_OE is True:
            for ind, text in enumerate(self.element_names_list):
                if ind != len(self.element_names_list):
                    posy = -self.lim
                    posx = xtick_pos[ind]

                if ind < len(self.element_names_list) - 1:
                    if (xtick_pos[ind + 1] - xtick_pos[ind]) < 1:
                        posy += 3
                ax1.text(posx, posy, text, rotation=60)
                ax2.text(posx, posy, text, rotation=60)

        # remove xticks and labels which are too close to eaeachhc other
        for ind in range(xtick_pos.shape[0] - 2, 1, -1):
            if ind > 0:  # < xtick_pos.shape[0]-1:
                if np.abs(xtick_pos[ind] - xtick_pos[ind + 1]) < 1:
                    xtick_label = np.delete(xtick_label, ind)
                    xtick_pos = np.delete(xtick_pos, ind)

        if debug:
            print("DEBUG:: plot_length", plot_length)
            print("DEBUG:: self.distance_list", self.distance_list)
            print("DEBUG:: beamline_length", beamline_length)
            print("DEBUG:: beamline_el_pos", beamline_el_pos)
            print("DEBUG:: xtick_label", xtick_label)
            print("DEBUG:: xtick_pos", xtick_pos)
        ax1.set_xticks(xtick_pos)
        ax1.set_xticklabels(xtick_label, rotation=45, ha="right")
        ax2.set_xticks(xtick_pos)
        ax2.set_xticklabels(xtick_label, rotation=45, ha="right")

        if lim_top is not False:
            ax1.set_xlim(lim_top[0], lim_top[1])
            ax1.set_ylim(lim_top[2], lim_top[3])
        if lim_side is not False:
            ax2.set_xlim(lim_side[0], lim_side[1])
            ax2.set_ylim(lim_side[2], lim_side[3])
        if save_img is True:
            plt.savefig(os.path.join(self.directory, img_name + extension))
        plt.tight_layout()
        plt.close(0)
        if show_img is True:
            plt.show()
        self.count_fig += 1

        return self.xh, self.yh, x / 1000, y
