"""
Holds the class SegmentedDmCommand that is used to create a dict that will be sent to the
segmented DM hardware as a command.

This module can be used to create a command in the following way:
    command = SegmentedDmCommand(dm_config_id, apply_flat_map=True,
                                 filename_flat='repo-path/hardware/iris-ao/CustomFLAT.ini')
    command.read_initial_values(segment_ptt_values)
    command.update_one_segment(segment_index, segment_ptt_values)
    command.display()
    dm_object.apply_shape(command)

Note that a command can be initialized from zeros, None, a list of tuples giving (P,T,T) for each segment,
an .ini file, or a .PTT### file. The PoppySegmentedDmCommand class can be used to create a command from
input Zernike coefficients.
Additionally, a single segment can be changed using the SegmentedDmCommand.update_one_segment()
method.

"""
from configparser import NoOptionError
import json
import os

import astropy.units as u
import matplotlib.pyplot as plt
import numpy as np
import poppy

from catkit.catkit_types import MetaDataEntry
from catkit.config import CONFIG_INI
from catkit.hardware.iris_ao import util

class SegmentedAperture(poppy.dms.HexSegmentedDeformableMirror):
    """
    Create a segmented aperture with Poppy using the parameters for the testbed
    and segmented DM from the config.ini file.

    To create an aperture, with the config.ini file loaded, run:
        segmented_aperture_obj = SegmentedAperture()

    :param dm_config_id: str, name of the section in the config_ini file where information
                         regarding the segmented DM can be found.
    :param rotation: float, rotation angle of the hex segmented DM in deg
    :attribute outer_ring_corners: bool, whether or not the segmented aperture includes
                                   the corner segments on the outer-most ring. If True,
                                   corner segments are included. If False, they are not
    :attribute center_segment: bool, whether of not the segmented aperture includes the
                               center segment. If True, the center segment is included.
                               If False, it is not.
    :attribute flat_to_flat: float, physical distance from flat to flat side of a DM segment
                             in units of mm
    :attribute gap: float, physical distance between DM segments in units of microns
    :attribute number_segments_in_pupil: int, number of segments included in the aperture for active
                                         control as specified by the "active_number_of_segments"
                                         parameter in the config file
    """

    def __init__(self, dm_config_id, rotation=0):
        # Set config sections
        self.dm_config_id = dm_config_id

        # Parameters specific to testbed setup being used
        self.rotation = rotation

        # Parameters specifc to the aperture and segmented DM being used
        self.outer_ring_corners = CONFIG_INI.getboolean(self.dm_config_id, 'include_outer_ring_corners')
        self.center_segment = CONFIG_INI.getboolean(self.dm_config_id, 'include_center_segment')
        self.flat_to_flat = CONFIG_INI.getfloat(self.dm_config_id, 'flat_to_flat_mm') * u.mm
        self.gap = CONFIG_INI.getfloat(self.dm_config_id, 'gap_um') * u.micron
        self.number_segments_in_pupil = CONFIG_INI.getint(self.dm_config_id, 'active_number_of_segments')

        # Get the specific segments
        self._num_rings = self.get_number_of_rings_in_pupil()
        self._segment_list = self.get_segment_list()

        super().__init__(name='Segmented DM',
                         rings=self._num_rings,
                         flattoflat=self.flat_to_flat,
                         gap=self.gap,
                         segmentlist=self._segment_list,
                         rotation=self.rotation)

    def get_number_of_rings_in_pupil(self):
        """
        Get the number of rings based on the number of active segments specified in the
        config.ini file.

        :return: num_rings: int, the number of full rings of hexagonal segments in the aperture
        """
        active_segs_per_ring = self.get_active_number_of_segments_per_ring()

        try:
            num_rings = active_segs_per_ring.index(self.number_segments_in_pupil)
        except ValueError as error:
            raise ValueError("Invalid number of segments. Please check your config.ini file.") from error

        return num_rings

    def get_active_number_of_segments_per_ring(self, max_number_of_rings=7):
        """
        Given pupil specifics, returns a list of length equal to the maximum
        number of rings where each element is the total number of segments in a
        pupil with that number of rings.

        Given the use of the center segment or outer ring corner segments,
        return a list of the number of total active segments in a pupil, of a variety
        of number of rings that is indicated by its index in the list. For example,
        the first element in the list, with index of 1, will give the total number of
        active segments in a pupil with 1 ring which include the center segment if that
        segment is active. Note: The maximum number of rings for the PTT489 from IrisAO is 7

        return: list of total number of active segments in a pupil for pupils with
                a number of rings as indicated by the index of in the list
        """
        max_number_segments_in_pupil_per_ring = self.get_max_number_segments_in_pupil_per_ring(max_number_of_rings)

        # If no outer corners, you will have 6 fewer segments in that outer ring
        if not self.outer_ring_corners:
            active_segments_in_pupil_per_ring = [num-6 if num > 6 else num for num in max_number_segments_in_pupil_per_ring]
        # If no center segment, you will have 1 fewer segment overall
        if not self.center_segment:
            active_segments_in_pupil_per_ring = [num-1 for num in max_number_segments_in_pupil_per_ring]

        if self.outer_ring_corners and self.center_segment:
            active_segments_in_pupil_per_ring = max_number_segments_in_pupil_per_ring

        return active_segments_in_pupil_per_ring

    def get_max_number_segments_in_pupil_per_ring(self, number_of_rings=7):
        """Returns a list of length equal to the maximum number of rings in a pupil
        with hexagonal segments where each element is the total number of segments in a
        pupil with that number of rings.

        :return: list, list of length equal to the total number of rings where each element
        is the total number of segments in a pupil with that number of rings.
        """
        max_number_segments_in_pupil_per_ring = [1,] # number of segments in the pupil per ring
        for i in np.arange(number_of_rings)+1:
            max_number_segments_in_pupil_per_ring.append(max_number_segments_in_pupil_per_ring[i-1]+i*6)

        return max_number_segments_in_pupil_per_ring

    def get_segment_list(self):
        """
        Grab the list of segments to be used in your pupil taking into account if your
        aperture has a center segment and/or the segments in the corners of the outer ring.
        This list is passed to Poppy to help create the aperture.

        :param num_rings: int, The number of rings in your pupil
        :return: list, the list of segments as passed into poppy
        """
        num_segs = self.total_number_segments_in_aperture(self._num_rings)
        seglist = np.arange(num_segs)

        # If no outer corners, you will have 6 fewer segments in that outer ring
        if not self.outer_ring_corners:
            inner_segs = seglist[:(num_segs-6*self._num_rings)]
            outer_segs = seglist[(num_segs-6*self._num_rings):]
            outer_segs = np.delete(outer_segs, np.arange(0, outer_segs.size,
                                                         self._num_rings)) # delete corner segs
            seglist = np.concatenate((inner_segs, outer_segs))

        if not self.center_segment:
            seglist = seglist[1:]
        return seglist

    def total_number_segments_in_aperture(self, number_of_rings=7):
        """
        For a segmented aperture of rings = number_of_rings, give the total number of
        segments in the aperture given an aperture of 1, 2, 3, etc. rings. Will return
        a list of the total number of segments in the *aperture* per ring where the ring
        is indicated by the index in the list.

        param: number_of_rings, int, the number of rings in the aperture
        return: a list of segments where the index of the value corresponds to the
                 number of rings. For example: index 0 corresponds with the center
                 where there is one segment.
        """
        number_segments_in_pupil_per_ring = self.get_max_number_segments_in_pupil_per_ring(number_of_rings)

        return number_segments_in_pupil_per_ring[number_of_rings]

class SegmentedDmCommand(object):
    """
    Handle segmented DM specific commands in terms of piston, tip and tilt (PTT) for
    each segment. Creates a dictionary of the form {seg: (piston, tip, tilt)} that can
    be loaded onto the hardware.

    Units of the loaded command are defined with the atrribute "dm_command_units" and are
    usually in um (for piston) and mrad (for tip and tilt).

    This class does NOT interact with hardware directly.

    :param dm_config_id: str, name of the section in the config_ini file where information
                         regarding the segmented DM can be found.
    :param apply_flat_map: If true, add the custom flat map correction to the data before creating the command
    :param filename_flat: string, full path to custom flat map, only needed if apply_flat_map=True
    :param rotation: float, rotation angle of the hex segmented DM in deg
    :attribute data: list of tuples, input data that can then be updated. This attribute never
                     never includes the custom flat map values; units are determined with the attribute dm_command_units
    :attribute apply_flat_map: bool, whether or not to apply the custom flat map
    :attribute source_pupil_numbering: list, numbering native to data
    :attribute command: dict, final command with flat if apply_flat_map = True; units are determined with the attribute dm_command_units
    :attribute filename_flat: str, full path to custom flat, only needed if apply_flat_map=True
    :attribute total_number_segments: int, total number of segments in DM, includes dead segments
    :attribute active_segment_list: int, number of active segments in the DM
    :attribute dm_command_units: tuple of floats, the units of the piston, tip, tilt
                                 values on the hardware
    :attribute aperture: poppy.dms.HexSegmentedDeformableMirror object, the aperture
                         that you are defining
    """

    def __init__(self, dm_config_id, apply_flat_map=False, filename_flat=None, rotation=0):
        self.dm_config_id = dm_config_id

        # Initialize class used to model the segmented aperture geometry
        self.aperture = SegmentedAperture(dm_config_id, rotation=rotation)

        # Determine if the custom flat map will be applied
        self.apply_flat_map = apply_flat_map
        if self.apply_flat_map:
            if filename_flat is None:
                raise ValueError("Must provide a filename for the segmented DM flat file, not None.")
            self.filename_flat = filename_flat
            # Check that the file exists
            if not os.path.isfile(self.filename_flat):
                raise FileNotFoundError(f"{self.filename_flat} either does not exists or is not currently accessible")

        # Establish segment information
        try:
            self.segments_in_pupil = json.loads(CONFIG_INI.get(self.dm_config_id, 'active_segment_list'))
            if len(self.segments_in_pupil) != self.aperture.number_segments_in_pupil:
                raise ValueError("The length of active_segment_list does not match the active_number_of_segments in the config.ini. Please update your config.ini.")
        except NoOptionError:
            self.segments_in_pupil = util.iris_pupil_naming(self.dm_config_id)

        # Set units for piston, tip, tilt
        dm_command_units = CONFIG_INI.get(self.dm_config_id, 'dm_ptt_units').split(',')
        self.dm_command_units = [u.Unit(dm_command_units[0]), u.Unit(dm_command_units[1]),
                                 u.Unit(dm_command_units[2])]

        # Initalize command
        self.data = util.create_zero_list(self.aperture.number_segments_in_pupil)
        self.input_data = None
        self.command = None

    def read_initial_command(self, segment_values):
        """
        Read a new command and assign to the attribute 'data'

        :param segment_values: str, list. Can be .PTT111, .ini files or a list with piston, tip,
                               tilt values in a tuple for each segment. See the load_command doc
                               string for more information.
        """
        self.input_data = self.data = self._read_command(segment_values)

    def _read_command(self, segment_values):
        """
        Read and return command from one of the allowed formats

        :param segment_values: str, list. Can be .PTT111, .ini files or a list with piston, tip,
                               tilt values in a tuple for each segment. See the load_command doc
                               string for more information.
        :return: list of tuples for each commanded segment
        """
        ptt_list, segment_names = util.read_segment_values(segment_values, self.dm_config_id)
        if segment_names is not None:
            command_list = []
            for seg_name in self.segments_in_pupil:  # Pull out only segments in the pupil
                ind = np.where(np.asarray(segment_names) == seg_name)[0][0]
                command_list.append(ptt_list[ind])
        else:
            command_list = ptt_list
        return command_list

    def get_data(self):
        """ Grab the current shape to be applied to the DM (does NOT include the custom flat map)
        """
        return self.data

    def to_command(self):
        """ Output command suitable (in the form a dictionary with an entry for each segment)
        for sending to the hardware driver. The custom flat map will be added only at this stage.
        """
        if self.apply_flat_map:
            command_data = self.add_map(self.filename_flat, return_new_map=True)
        else:
            command_data = self.data
        command_dict = dict(zip(self.segments_in_pupil, command_data))

        return command_dict

    def update_one_segment(self, segment_ind, ptt_tuple, add_to_current=True):
        """ Update the value of one segment by supplying the new command that will be added
        to or will replace the current PTT tuple on this segment only. To identify the segment to be changed, give
        the *command index* (not its name like in the GUI) of that segment in the active segment list. This will be added
        to the current value only if the add_to_current flag is set to True, otherwise, value
        given will replace the current value.

        :param segment_ind: int, for the segment in the pupil that is to be updated,
                            provide the index of it's location in the active segment list
        :param ptt_tuple: tuple with three values for piston, tip, and tilt in um and mrad unless specified otherwise
                          with self.dm_command_units
        """
        if add_to_current:
            command_list = util.create_zero_list(self.aperture.number_segments_in_pupil)
            command_list[segment_ind] = ptt_tuple
            self.add_map(command_list)
        else:
            self.data[segment_ind] = ptt_tuple

    def add_map(self, segment_values_to_add, return_new_map=False):
        """
        Add a command to the one already loaded.

        :param new_command: str or list (.PTT111 or .ini file, or list, for example, from POPPY)
        """
        original_data = self.get_data()
        data_to_add = self._read_command(segment_values_to_add)

        new_map = [tuple(map(sum, zip(orig, addition))) for orig,
                   addition in zip(original_data, data_to_add)]

        if return_new_map:
            return new_map
        else:
            self.data = new_map

    def to_ini(self, filename, out_dir=''):
        """
        Write the command to a .ini file. Note: This will NOT include the applied custom flat
        map values
        :param filename: str, name of ini file to be written out
        :param out_dir: str, name of directory where ini file will be saved
        """
        if not filename.endswith(".ini"):
            filename += ".ini"

        data = self.get_data()
        command = dict(zip(self.segments_in_pupil, data))
        path = os.path.join(out_dir, filename)
        util.write_ini(command, path, dm_config_id=self.dm_config_id)

    def get_extra_meta_data(self):
        """
        Create meta data to be saved with fits files that gives the ptt values/segment
        and if the flat map was applied.The values saved will NOT include custom flat map PTT values.
        """
        metadata = []
        if self.apply_flat_map:
            metadata.append(MetaDataEntry("Flat Name", "FLATMAP", self.filename_flat,
                                          "Flat map name/file if applied"))
        rounded_ptt_list = round_ptt_list(self.data, decimals=4)    # since self.data is usually in units of (um, mrad, mrad), it is ok to round to 4 digits here)
        for seg, ptt in zip(self.segments_in_pupil, rounded_ptt_list):
            metadata.append(MetaDataEntry(f"Segment {seg}", f"SEG{seg}", str(ptt),
                                          f"Piston/GradX/GradY applied for segment {seg}"))
        return metadata

    def display(self, display_wavefront=True, display_psf=True, psf_rotation_angle=0., vmax_opd=0.5e-6*u.meter, vmin_psf=1e-8, vmax_psf=1e-2,
                save_figures=True, figure_name_prefix='', out_dir=''):
        """
        Display either the deployed mirror state ("wavefront") and/or the PSF created
        by this mirror state.

        :param display_wavefront: bool, If true, display the deployed mirror state
        :param display_psf: bool, If true, display the simulated PSF created by the
            mirror state
        :param psf_rotation_angle: int/float, Degree value by which to rotate the simulated
            PSF in order to match your output
        :param figure_name_prefix: str, String to be added to filenames of output figures
        :param out_dir: str, name of output directory
        :param save_figures: bool, If true, save out the figures in the directory specified
                             by out_dir
        """
        # update opd on mirror
        self.apply_current_wavefront()

        # make figure
        if figure_name_prefix:
            figure_name_prefix = f'{figure_name_prefix}_'
        if display_wavefront:
            self.plot_wavefront(figure_name_prefix, out_dir, vmax=vmax_opd, save_figure=save_figures)
        if display_psf:
            self.plot_psf(rotation_angle=psf_rotation_angle, vmin=vmin_psf, vmax=vmax_psf, figure_name_prefix=figure_name_prefix, out_dir=out_dir,
                          save_figure=save_figures)

    def plot_wavefront(self, figure_name_prefix, out_dir, vmax=0.5e-6*u.meter, save_figure=True):
        """
        Plot the deployed mirror state (wavefront error)

        :param figure_name_prefix: str, String to be added to filenames of output figures
        :param out_dir: str, name of output directory
        :param vmax: astropy unit quantity, the maximum value to display in the plot (as expected
                     by Poppy)
        :param save_figure: bool, If true, save out the figures in the directory specified by
                            out_dir
        """
        plt.figure()
        self.aperture.display(what='opd', title='Wavefront error applied to the active segments',
                              opd_vmax=vmax)
        if save_figure:
            plt.savefig(os.path.join(out_dir, f'{figure_name_prefix}wfe_on_dm.png'))
            plt.close()
        else:
            plt.show()

    def apply_current_wavefront(self):
        """
        Get the deployed mirror state (wavefront error)
        To make plot:
        total_iris = iris_ao.HicatSegmentedDmCommand()
        total_iris.read_initial_command([tuple(row) for row in command_to_plot])
        total_iris.apply_current_wavefront()
        command = total_iris.aperture.sample(what='opd')
        image = ax.matshow(command)
        """
        display_data = set_to_dm_limits(self.data)
        # Convert the PTT list from DM to Poppy units
        converted_list = convert_ptt_units(display_data, tip_factor=1, tilt_factor=-1,
                                           starting_units=self.dm_command_units,
                                           ending_units=(u.m, u.rad, u.rad))
        # We want to round to four significant digits when in DM units (um, mrad, mrad).
        # Here, we are in SI units (m, rad, rad), so we round to the equivalent, 10 decimals.
        rounded_list = round_ptt_list(converted_list, decimals=10)
        for seg, values in zip(self.aperture.segmentlist, rounded_list):
            self.aperture.set_actuator(seg, values[0], values[1], values[2])

    @poppy.utils.quantity_input(display_wavelength=u.nm)   # decorator provides a check on input units
    def plot_psf(self, wavelength=640*u.nm, figure_name_prefix=None, rotation_angle=0,  out_dir=None, vmin=1e-8,
                 pixelscale=0.010, instrument_fov=1.0,
                 vmax=1e-2, save_figure=True):
        """
        Plot the simulated PSF based on the mirror state. Optionally save figure to a file as well

        :param wavelength: Wavelength to calculate and display a PSF for
        :param pixelscale: pixel scale to calculate and display the PSF on
        :param instrument_fov: instrument field of view to calculate and display the PSF on
        :param rotation_angle: float, the rotation to apply to the PSF image in order to match
                               our testbed. Note that this value is specific to each testbed
        :param save_figure: bool, If true, save out the figures in the directory specified by
                            out_dir
        :param figure_name_prefix: str, String to be added to filenames of output figures
        :param out_dir: str, name of output directory
        :param vmin: float, the minimum value to display in the plot
        :param vmax: float, the maximum value to display in the plot
       """
        plt.figure()
        osys = poppy.OpticalSystem()
        osys.add_pupil(self.aperture)
        osys.add_detector(pixelscale=pixelscale, fov_arcsec=instrument_fov)
        osys.add_rotation(angle=rotation_angle)

        psf = osys.calc_psf(wavelength=wavelength)
        poppy.display_psf(psf, vmin=vmin, vmax=vmax,
                          title='PSF created by the shape put on the active segments')
        if save_figure:
            if out_dir is None or figure_name_prefix is None:
                raise ValueError("Must specify out_dir and figure_name_prefix if save_figure is True")
            plt.savefig(os.path.join(out_dir, f'{figure_name_prefix}simulated_psf.png'))
            plt.close()

def load_command(segment_values, dm_config_id,
                 apply_flat_map=True, filename_flat=None):
    """
    Loads the segment_values from a file or list and returns a SegmentedDmCommand object.

    There are only two allowed segment mappings for the input formats, Native and Centered Pupil.
    See the README for the Iris AO for more details.

    :param segment_values: str, list. Can be .PTT111, .ini files or a list with piston, tip,
                           tilt values in a tuple for each segment. For the list, the first
                           element is the center or top of the innermost ring of the pupil,
                           and subsequent elements continue up and/or clockwise around the
                           pupil (see README for more information)
    :param dm_config_id: str, name of the section in the config_ini file where information
                         regarding the segmented DM can be found.
    :param apply_flat_map: bool, whether to apply a flat map in addition to the data when sending command to hardware
    :param filename_flat: string, full path to the custom flat map, only needed if apply_flat_map=True
    :return: SegmentedDmCommand object representing the command dictionary.
    """
    dm_command_obj = SegmentedDmCommand(dm_config_id=dm_config_id,
                                        apply_flat_map=apply_flat_map,
                                        filename_flat=filename_flat)
    dm_command_obj.read_initial_command(segment_values)

    return dm_command_obj

## POPPY
def round_ptt_list(ptt_list, decimals=3):
    """
    Make sure that the PTT coefficients are rounded to a reasonable number of decimals

    :param ptt_list: list, of tuples existing of piston, tip, tilt, values for each
                     segment in a pupil
    :return: list of coefficients for piston, tip, and tilt, for your pupil rounded to
             the specified number of decimal places
    """
    return [(np.round(ptt[0], decimals), np.round(ptt[1], decimals),
             np.round(ptt[2], decimals)) for ptt in ptt_list]

def convert_ptt_units(ptt_list, tip_factor, tilt_factor, starting_units, ending_units):
    # TODO: this should follow the logic from here: https://github.com/spacetelescope/catkit/blob/8c1e22e8f96d51f4453bba27c4710d27da5d37c5/catkit/emulators/iris_ao_controller.py#L50
    # This is used in display, where it's not super important, and it's used in the global Zernike commands, so those
    # are not super reliable at this point.
    """
    Convert the PTT list to or from Poppy units and the segmented DM units.

    Note that this has been created for the IrisAO segmented DMs, therefore
    the tip and tilt values are swapped with what Poppy desginates.

    Poppy to segmented DM:
    - tip_factor = -1
    - tilt_facotr = 1
    - starting_units = (u.m, u.rad, u.rad)
    - ending_units = dm_ptt_units from the config.ini file

    Segmented DM to Poppy
    - tip_factor = 1
    - tilt_factor = -1
    - starting_units = dm_ptt_units from the config.ini file
    - ending_units = (u.m, u.rad, u.rad)

    :param ppt_list: list, of tuples existing of piston, tip, tilt, values for each
                     segment in a pupil, in the respective starting_units
    :param tip_factor: int, either -1 or 1 based on the information above
    :param tilt_factor: int, either -1 or 1 based on the information above
    :param starting_units: tuple or list of the units associated with the piston, tip,
                           tilt values respectively of the input ptt_list
    :param ending_units: tuple_or_list of the units associated with the piston, tip,
                           tilt values respectively of the expected output
    :return: list of tuples of the piston, tip, tilt values for each segment listed,
             in the respective ending_units
    """
    converted = [(ptt[0]*(starting_units[0]).to(ending_units[0]),
                  tip_factor*ptt[2]*(starting_units[2]).to(ending_units[2]),
                  tilt_factor*ptt[1]*(starting_units[1]).to(ending_units[1])) for ptt in ptt_list]
    return converted

def set_to_dm_limits(ptt_list, limit=5.):
    """
    Check that the values for piston, tip, and tilt are not exceeding the hardware
    limit and reset to limit if limit is exceeded. These limits are the same as what
    the IrisAO GUI has set.

    :param ppt_list: list, of tuples existing of piston, tip, tilt, values for each
                     segment in a pupil, in DM units
    :param limit: float, in DM units. Default = 5.
    :return: list of tuples of the piston, tip, tilt values in DM units for each segment listed
             such that none of the values exceed the limit
    """
    updated = [tuple(min(i, limit) for i in ptt) for ptt in ptt_list]

    return updated

def get_wavefront_from_coeffs(coeff_list, basis):
    """
    Get the wavefront from the coefficients created by the basis given. This gives
    the per-segment wavefront based on the global coefficients given and the basis
    created.

    :params basis: POPPY Segment_PTT_Basis object, basis created that represents pupil
    :params coeff_list: list, per-segment list of piston, tip, tilt values for
                         the pupil described by basis
    :return wavefront, POPPY wavefront
    """
    wavefront = poppy.zernike.opd_from_zernikes(coeff_list, basis=basis)
    return wavefront

class PoppySegmentedDmCommand(SegmentedDmCommand):
    """ Create segmented DM command based on provided wavefront parameters

    This is currently limited to global shapes defined by Zernike coefficients

    The output is
    a SegmentedDmCommand containing piston, tip, tilt values with DM units,
    usually (um, mrad, mrad), respectively, for each segment.

    This class inherits the SegmentedAperture class.

    To use to get command for the segmented DM:
      command = PoppySegmentedDmCommand(coefficients, dm_config_id)
      dm_object.apply_shape(command)

    The method .to_dm_list() will place the command in the correct units to be sent to the hardware;
    this happens automatically in __init__.

    :param global_coefficients: list of global zernike coefficients in the form
                                [piston, tip, tilt, defocus, ...] (Noll convention)
                                in meters of optical path difference (not waves)
    :param dm_config_id: str, name of the section in the config_ini file where information
                         regarding the segmented DM can be found.
    :param display_wavelength: float, wavelength in nm of the poppy optical system used for
                        (extremely oversimplified) focal plane simulations
    :param rotation: float, rotation angle of the hex segmented DM in deg
    :attribute radius: float, half of the flat-to-flat distance of each segment
    :attribute num_terms: int, total number of PTT values on all segments (= 3 x number of segments)
    :attribute dm_command_units: list, the units of the piston, tip, tilt (respecitvely)
                                 values when coming from the DM or DM command
    :attribute global_coefficients: list of global zernike coefficients in the form
                                [piston, tip, tilt, defocus, ...] (Noll convention)
                                in meters of optical path difference (not waves)
    :attribute basis: poppy.zernike.Segment_PTT_Basis object, basis based on the characteristics
                      of the segmented DM being used
    :attribute list of coefficients: list of piston, tip, tilt coefficients in units of m, rad, rad
    """
    @poppy.utils.quantity_input(display_wavelength=u.nm)   # decorator provides a check on input units
    def __init__(self, global_coefficients, dm_config_id, rotation=0, **kwargs):
        # Initialize parent class
        super().__init__(dm_config_id,rotation=rotation, **kwargs)

        self.radius = (self.aperture.flat_to_flat/2).to(u.m)
        self.num_terms = (self.aperture.number_segments_in_pupil) * 3

        # Grab the units of the DM for the piston, tip, tilt values to convert to
        dm_command_units = CONFIG_INI.get(self.dm_config_id, 'dm_ptt_units').split(',')
        self.dm_command_units = [u.Unit(dm_command_units[0]), u.Unit(dm_command_units[1]),
                                 u.Unit(dm_command_units[2])]

        self.global_coefficients = global_coefficients

        # Create the specific basis for this pupil
        self.basis = self.create_ptt_basis()

        # Create list of segment coefficients
        self.list_of_coefficients = self.get_list_from_global()

        # Set the command data based on that list of coefficients
        self.read_initial_command(self.to_dm_list())

    def create_ptt_basis(self):
        """
        Create the basis needed for getting the per/segment coeffs back

        :return: Poppy Segment_PTT_Basis object for the specified pupil
        """
        pttbasis = poppy.zernike.Segment_PTT_Basis(rings=self.aperture._num_rings,
                                                   flattoflat=self.aperture.flat_to_flat,
                                                   gap=self.aperture.gap,
                                                   segmentlist=self.aperture._segment_list)
        return pttbasis

    def create_wavefront_from_global(self, global_coefficients, wavelength=640*u.nm):
        """
        Given an list of global coefficients, create wavefront

        :param global_coefficients: list of global zernike coefficients in the form
                                    [piston, tip, tilt, defocus, ...] (Noll convention)
                                    in meters of optical path difference (not waves)
        :param wavelength: Wavelength to create this wavefront at; does not actually affect the
                                    values of the OPD, which is modeled as wavelength-independent.
        :return: Poppy ZernikeWFE object, the global wavefront described by the input coefficients
        """
        wavefront = poppy.ZernikeWFE(radius=self.radius, coefficients=global_coefficients)
        wavefront_out = wavefront.sample(wavelength=wavelength,
                                         grid_size=2*self.radius,
                                         npix=512, what='opd')
        return wavefront_out

    def get_coeffs_from_pttbasis(self, wavefront):
        """
        From a the speficic pttbasis, get back the coeff_list that will be sent as
        a command to the Iris AO

        :param wavefront: Poppy ZernikeWFE object, the global wavefront over the
                          pupil
        :return: list, (piston, tip, tilt) values for each segment in the pupil
                 in units of [m] for piston, and [rad] for tip and tilt
        """
        coeff_list = poppy.zernike.opd_expand_segments(wavefront, nterms=self.num_terms,
                                                        basis=self.basis)
        coeff_list = np.reshape(coeff_list, (self.aperture.number_segments_in_pupil, 3))
        return coeff_list

    def get_list_from_global(self):
        """
        From a global coeff list, get back the per-segment coefficient list

        :return: list of coefficients for piston, tip, and tilt, for your pupil
        """
        wavefront = self.create_wavefront_from_global(self.global_coefficients)
        coeffs_list = self.get_coeffs_from_pttbasis(wavefront)

        return coeffs_list

    def to_dm_list(self):
        """
        Convert the PTT list to DM hardware units so that it can be passed to the
        SegmentDmCommand class

        :return: list of coefficients for piston, tip, and tilt, for your pupil, in DM units,
                 usually um, mrad, mrad unless otherwise specified in the config
        """
        input_list = self.list_of_coefficients
        # Convert from Poppy's m, rad, rad to the DM units
        input_list = convert_ptt_units(input_list, tip_factor=-1, tilt_factor=1,
                                       starting_units=(u.m, u.rad, u.rad),
                                       ending_units=self.dm_command_units)
        coeffs_list = round_ptt_list(input_list)    # since input_list is usually in units of (um, mrad, mrad), it is ok to round to 4 digits here)

        return coeffs_list


#usage of the tests:
# ```
#   from functions import *
#   test_final('your__testfile_name')
# ```

#mongoDB setup
import pymongo
from pymongo import MongoClient
from bson.objectid import ObjectId
client = MongoClient()
client = MongoClient('localhost', 27017)

#3 list XML's
import glob
import json
def list_xmls(glob_paths,out_file=None):
    files = []
    for path in glob_paths:
        files.extend(glob.glob(path+"*.tei.xml"))
    print("{} files".format(len(files)))
    for i,f in enumerate(files):
        #replace the "\" returned in the glob paths
        files[i] = f.replace('\\','/')

    files.reverse()
    caches = set()
    results = []
    #remove duplicates
    for item in files:
        prefix = item.rsplit('v', 2)[0]
        # check whether prefix already exists
        if prefix not in caches:
            results.append(item)
            caches.add(prefix)
    results.reverse()

    return(results)

#4 create fulltexts array
from contextlib import suppress
import grobid_tei_xml
from os import path

def create_texts(files, out_file=None):
    texts = []
    for file_id,file_path in enumerate(files):

        title = None
        authors = None
        text = None
        doi = None
        citation_count = None
        abstract = None

        with suppress(Exception):
            with open(file_path, 'r', encoding="utf8") as d_file:
                doc = grobid_tei_xml.parse_document_xml(d_file.read())
        with suppress(Exception):
            text = doc.body
        with suppress(Exception):
            title= doc.header.title
        with suppress(Exception):
            authors= ", ".join([a.full_name for a in doc.header.authors])
        with suppress(Exception):
            doi= str(doc.header.doi)
        with suppress(Exception):
            citation_count= str(len(doc.citations))
        with suppress(Exception):
            abstract= doc.abstract

        texts.append({
            "ID":file_id,
            "path":file_path,
            "text":text,
            "title":title,
            "authors":authors,
            "doi":doi,
            "citation_count":citation_count,
            "abstract":abstract
        })
        print("{} files without duplicates".format(len(texts)),end='\r')
    print('')
    if (out_file != None):
        with open('./'+out_file+'.json', 'w') as fout:
            json.dump(texts, fout)
    return(texts)

from nltk.tokenize import sent_tokenize
def create_phrses(texts,phrases_db_name):
    phr_db = client[phrases_db_name]
    db_all_phrases = phr_db.all_phrases
    all_phrases_on_db = db_all_phrases.count_documents({})
    all_phrases = []
    cont_file = 0
    for text in texts:
        cont_file = cont_file+1
        for section in text:
            if (section in ["text","title","abstract"]):
                if (text[section]!= None):
                    #print(text[section])
                    for phrase in sent_tokenize(text[section]):
                        phr = {
                            "_id":len(all_phrases),
                            "path":text["path"],
                            "phrase":phrase.lower(),
                            "lenght":len(phrase),
                            "section":section
                        }
                        phrase_id = db_all_phrases.insert_one(phr)
                        all_phrases.append(phrase_id)
                        print("file: {} ----,{} phrases in total".format(cont_file,len(all_phrases)),end='\r')
    return(all_phrases)

#search for texts containing any strings in a set
from collections import namedtuple
from langdetect import detect
import langdetect.lang_detect_exception as ldError

def search_texts(search_set,fields,db_name,phrases_db_name):
    db = client[db_name]
    phr_db = client[phrases_db_name]
    phrases = db.phrases
    phrases_on_db = phrases.count_documents({})
    if (phrases_on_db == 0):
        print("The database is empty. running your query...")
        matches = run_query(search_set,fields,db_name,phrases_db_name);
    else:
        valid_ans = bool(False)
        while (not valid_ans):
            redo = str(input("Database already has {} phrases, do you want to retry query? Y-n".format(phrases_on_db)))
            if (redo in {"Y","y"}):
                matches =run_query(search_set,fields,db_name,phrases_db_name);
                valid_ans = True
            elif(redo in {"N","n"}):
                cursor = phrases.find({})
                matches = [phrase['phrase'] for phrase in cursor]
                valid_ans = True
            else:
                print("Please answer with Y,y(Yes) or N,n(No)...")
    print('')
    return(matches)

def run_query(search_set,sections,db_name,phrases_db_name):
    db = client[db_name]
    phrases = db.phrases
    phrases.drop()
    phr_db = client[phrases_db_name]
    all_phrases = phr_db.all_phrases
    matches = []
    for phrase in all_phrases.find():
        if (phrase["section"] in sections):
            match_words = []
            for word in search_set:
                if (re.search(r"\b" + re.escape(word) + r"\b", phrase["phrase"].lower())):
                    if(len(match_words)>0):
                        phrase_id = phrases.update_one({"_id" :int(phrase["_id"])},{"$set" : {"match_word":match_words}})
                    else:
                        remove_non_alpha = re.compile('[^a-zA-Z \-]')
                        ratio =  len(remove_non_alpha.sub('',phrase["phrase"].lower()))/len(phrase["phrase"].lower())
                        if (ratio>0.8):
                            try:
                                lang = detect(phrase["phrase"])
                                if (lang=='en'):
                                    match = phrase
                                    match["a_id"] = match["_id"]
                                    match["_id"] = len(matches)
                                    match["match_word"] = [word]
                                    phrase_id = phrases.insert_one(match)
                                    #matches.append(Phrase(path=text["path"],phrase=phrase.lower(),lenght=len(phrase)))
                                    matches.append(phrase_id)
                                    print("{} phrase matches".format(len(matches)),end='\r')
                            except ldError.LangDetectException:
                                pass
                    match_words.append(word)
    return(matches)

#5 remove stopwords & clean text
from sklearn.feature_extraction.text import CountVectorizer
import pandas as pd
import numpy as np
import nltk
from nltk.corpus import stopwords
from nltk.stem.porter import PorterStemmer

nltk.download('stopwords')
nltk.download('punkt')
from nltk.tokenize import word_tokenize
import re
remove_non_alpha = re.compile('[^a-zA-Z \-]')

def clean_text(db_name,remove_w=[],join_w=[],stem=False):
    db = client[db_name]
    phrases = db.phrases
    clean_phrases = []
    cursor = phrases.find({})
    rm_w_count =0
    join_w_count =0
    for n,phrase in enumerate(cursor):
        #get phrase text (sentence)
        sentence = phrase["phrase"]

        #remove words from a list.
        for rm_word in remove_w:
            while (re.search(r"\b" + re.escape(rm_word) + r"\b", sentence)):
                rm_w_count = rm_w_count+1
                sentence = re.sub(r"\b" + re.escape(rm_word) + r"\b","", sentence)

        #join words from a list.
        for join_word in join_w:
            while (re.search(r"\b" + re.escape(join_word) + r"\b", sentence)):
                join_w_count = join_w_count+1
                hyphen_word = join_word.replace(" ", "-")
                if(join_word == hyphen_word):
                    break
                sentence = re.sub(r"\b" + re.escape(join_word) + r"\b",hyphen_word, sentence)

        #remove non_alphanumeric characters
        sentence = remove_non_alpha.sub('',sentence.lower())

        #tokenize sentence
        tokens = word_tokenize(sentence)

        #remove Punctuation (already done)
        #words = [word for word in tokens if word.isalpha()]

        #remove stopwords
        stop_words = set(stopwords.words('english'))
        words = [w for w in tokens if not w in stop_words]

        #stem words
        if(stem == True):
            from nltk.stem.porter import PorterStemmer
            porter = PorterStemmer()
            words = [porter.stem(word) for word in tokens]

        #un-tokenize string
        cleant = " ".join(list(words))

        #append clean string to list
        clean_phrases.append(cleant)
        print(str(n)+"-"+str(len(tokens)),end='\r')
    print("{} palavras removidas".format(rm_w_count))
    print("{} palavras unidas".format(join_w_count))
    return(clean_phrases)

def clean_all_text(db_name,remove_w=[],join_w=[],stem=False):
    db = client[db_name]
    phrases = db.all_phrases
    clean_phrases = []
    cursor = phrases.find({})
    rm_w_count =0
    join_w_count =0
    for n,phrase in enumerate(cursor):
        remove_non_alpha = re.compile('[^a-zA-Z \-]')
        ratio =  len(remove_non_alpha.sub('',phrase["phrase"].lower()))/len(phrase["phrase"].lower())
        if (ratio>0.8):
            #get phrase text (sentence)
            sentence = phrase["phrase"]

            #remove words from a list.
            for rm_word in remove_w:
                while (re.search(r"\b" + re.escape(rm_word) + r"\b", sentence)):
                    rm_w_count = rm_w_count+1
                    sentence = re.sub(r"\b" + re.escape(rm_word) + r"\b","", sentence)

            #join words from a list.
            for join_word in join_w:
                while (re.search(r"\b" + re.escape(join_word) + r"\b", sentence)):
                    join_w_count = join_w_count+1
                    hyphen_word = join_word.replace(" ", "-")
                    if(join_word == hyphen_word):
                        break
                    sentence = re.sub(r"\b" + re.escape(join_word) + r"\b",hyphen_word, sentence)

            #remove non_alphanumeric characters
            sentence = remove_non_alpha.sub('',sentence.lower())

            #tokenize sentence
            tokens = word_tokenize(sentence)

            #remove Punctuation (already done)
            #words = [word for word in tokens if word.isalpha()]

            #remove stopwords
            stop_words = set(stopwords.words('english'))
            words = [w for w in tokens if not w in stop_words]

            #stem words
            if(stem == True):
                from nltk.stem.porter import PorterStemmer
                porter = PorterStemmer()
                words = [porter.stem(word) for word in tokens]

            #un-tokenize string
            cleant = " ".join(list(words))

            #append clean string to list
            clean_phrases.append(cleant)
            print(str(n)+"-"+str(len(tokens)),end='\r')
    print("")
    print("{} palavras removidas".format(rm_w_count))
    print("{} palavras unidas".format(join_w_count))
    return(clean_phrases)

#5 create n_gram and vocab
from sklearn.feature_extraction.text import CountVectorizer
def create_n_gram(phrases):
    doc = pd.DataFrame(phrases)
    doc.columns = ['text']
    vect = CountVectorizer()
    vects = vect.fit_transform(doc.text)
    return(vect,vects)

#clean n_gram (optional, DOES NOT WORK with porter stemmer)
def clean_n_gram():
    with open('./american-english-huge','r') as f:
            eng_words_huge=set()
            for line in f:
                    strip_lines=line.strip()
                    listli=strip_lines.split()
                    eng_words_huge.add(listli[0])

#6 run LDA
import lda

def run_lda(vect,vects,topic_number,db_name):
    db = client[db_name]
    phrases = db.phrases
    topics = db.topics
    X = vects.toarray()
    model = lda.LDA(n_topics=topic_number, n_iter=2500, random_state=1,refresh=10)
    model.fit(X)  # model.fit_transform(X) is also available
    #lista de palavras e probabilidades de cada tópico:
    n_words_on_topic = (model.nzw_ > 0).sum(axis=1)
    top_word_indices_of_topics = [(-model.nzw_[a]).argsort()[:n_words_on_topic[a]] for a in range(len(n_words_on_topic))]
    #save topics to the Db
    topics_on_db = topics.count_documents({})
    for i,top_word_indices in enumerate(top_word_indices_of_topics):
        w_p = [{"word":(vect.get_feature_names())[word],"prob":np.around(model.topic_word_,5)[i][word]} for word in top_word_indices]
        this_topic = {"_id":i,"word_probabilities":w_p}
        if (topics_on_db == 0):
            topics.insert_one(this_topic)
    #add topics to the phrases in the Db
    n_topics_on_document = (model.ndz_ > 0).sum(axis=1)
    top_topic_indices_of_documents = [(-model.ndz_[a]).argsort()[:n_topics_on_document[a]] for a in range(len(n_topics_on_document))]
    for i,top_topic_indices in enumerate(top_topic_indices_of_documents):
        doc_topics = [{"topic": int(topic),"prob":float(np.around(model.doc_topic_,5)[i][topic])} for topic in top_topic_indices]
        phrases.update_one({"_id" :int(i)},{"$set" : {"topics":doc_topics}})
    return(model)

#populate_DB
import datetime
import ast
import uuid
import timeit

def run_analysis():
    #setup mongo client
    client = MongoClient('localhost', 27017)
    #create/access history of
    hist_db = client.lda_arxiv_log
    history = hist_db.history
    entries = history.count_documents({})
    latest = history.find_one({"$query": {}, "$orderby": {"$natural" : -1}})
    if(latest == None):
        print("no Analysis was found in the database... \n we will have to create yours from scratch!")
        this_analysis = input_analysis()
    else:
        valid_ans = bool(False)
        while (not valid_ans):
            redo = input("Do you want to redo an analysis? Y-n")
            if(redo in {"Y","y"}):
                analysis_hist = [hi for hi in history.find({"$query": {}, "$orderby": {"$natural" : -1}})]
                nicknames_list = [hi["nick_name"] for hi in analysis_hist]
                print("these are all the analysis you made:")
                for hi in analysis_hist:
                    print(hi)
                valid_ans = bool(False)
                while (not valid_ans):
                    analysis_old_name = input("please give us the nickname of the analysis you want to redo:")
                    this_analysis = next((hi for hi in analysis_hist if hi["nick_name"] == analysis_old_name), None)
                    if(this_analysis != None):
                        valid_ans = True
                    else:
                        print("please write a valid nickname")
                valid_ans =True
            elif(redo in {"N","n",}):
                this_analysis = input_analysis()
                valid_ans =True
            else:
                print("Please answer with Y,y(Yes) or N,n(No)...")

    glob_paths_ex = ["./pdf/*/"]
    query_words_ex = {"artificial intelligence", "machine learning", "m.l.","a.i."}
    query_sections_ex = {"abstract","text"}
    all_phrases_db_name = "ALL_PHRASES_ARXIV"

    #CHECK IF ALL_PHRASES EXISTS
    phr_db = client[all_phrases_db_name]
    db_all_phrases = phr_db.all_phrases
    all_phrases_on_db = db_all_phrases.count_documents({})
    if(all_phrases_on_db == 0):
            print("###list_xmls###")
            files = list_xmls(this_analysis["docs_path_list"])
            print("###create_texts###")
            test_text = create_texts(files)
            print("###create_phrses###")
            all_phrases = create_phrses(test_text,all_phrases_db_name)
    else:
        cursor = db_all_phrases.find({})
        all_phrases = [phrase['_id'] for phrase in cursor]
        print("###{} phrses on DB###".format(len(all_phrases)))

    print("###search_texts###")
    test_phrases = search_texts(this_analysis["query"]["words"],this_analysis["query"]["sections"],this_analysis["db_name"],all_phrases_db_name)
    print("###clean_text###")
    test_clean = clean_text(this_analysis["db_name"],join_w=this_analysis["query"]["words"])
    print("###create_n_gram###")
    (vect,vects) = create_n_gram(test_clean)
    starttime = timeit.default_timer()
    print(">>> start time is :",starttime)
    model = run_lda(vect,vects,this_analysis["n_topics"],this_analysis["db_name"])
    print("<<< time difference is :", timeit.default_timer() - starttime)
    this_analysis["date"]= datetime.datetime.utcnow()
    this_analysis["_id"]= str(uuid.uuid4())
    history.insert_one(this_analysis)
    print("thanks, your analysis parameters are stored with the nickname: {}.".format(this_analysis["nick_name"]))

def input_analysis():
        hist_db = client.lda_arxiv_log
        history = hist_db.history
        this_analysis ={}
        print("The databadeses in your machine are:")
        dbs_in_machine = []
        for db in client.list_databases():
            print(db)
            dbs_in_machine.append(db["name"])
        database_name = "arxiv_LDA_{}".format(input("Write the name of the new (or old) database: arxiv_LDA_..."))
        this_analysis["db_name"] = database_name
        #get docs_path_list
        valid_ans = bool(False)
        while (not valid_ans):
            try:
                this_analysis["docs_path_list"] = ast.literal_eval(input('Write a list of paths where the documents are found: ["./in", "./this/way", "./pdf/*/"]') or '["./pdf/*/"]' )
                valid_ans = (type(this_analysis["docs_path_list"])==list)
                #check if paths exist...
            except:
                valid_ans = False
            if (valid_ans == False):print("please write your LIST of paths in pythonic syntax")

        this_analysis["query"] = {}
        valid_ans = bool(False)
        while (not valid_ans):
            try:
                this_analysis["query"]["words"] = ast.literal_eval(input('Write your set of query words: [ "in", "this way"]') or '["artificial intelligence", "machine learning", "m.l.","a.i."]')
                valid_ans = (type(this_analysis["query"]["words"])==list)
            except:
                valid_ans = False
            if (valid_ans == False):print("please write your LIST of words in pythonic syntax")

        valid_ans = bool(False)
        while (not valid_ans):
            try:
                this_analysis["query"]["sections"] = ast.literal_eval(input('Write your set of sections that you want to query: [ "in", "this way"]') or '["abstract","text"]')
                valid_ans = (type(this_analysis["query"]["sections"])==list)
            except:
                valid_ans = False
            if (valid_ans == False):print("please write your LIST of sections in pythonic syntax")

        valid_ans = bool(False)
        while (not valid_ans):
            try:
                this_analysis["n_topics"] = int(input("How many topics do you wish?") or 20)
                valid_ans = True
            except:
                valid_ans = False
            if (valid_ans == False):print("Number of topics must be an Integer")

        valid_ans = bool(False)
        nicknames_list = [hi["nick_name"] for hi in history.find({"$query": {}, "$orderby": {"$natural" : -1}})]
        while (not valid_ans):
            this_analysis["nick_name"] = input("we are almost there... Just give your analysis an nickname:")
            if (this_analysis["nick_name"] not in nicknames_list):
                valid_ans = True
            else:
                print("this nickname already exists please give it an Unique one!")
        return(this_analysis)


from dotenv import load_dotenv
from pipelines_service import AzurePipelineService
from utils import get_property

def main():
    load_dotenv()
    pipelines_service = AzurePipelineService()

    prop = get_prop_to_filter_by()
    value = get_prop_value_to_filter_by()

    pipelines = list(pipelines_service.get_pipelines_by_prop_with_initial_value_of(prop, value))
    pipelines_ids = get_property(pipelines, 'id')
    pipelines_names = get_property(pipelines, 'name')
    pipelines_names.sort()

    display_retrieved_pipelines(pipelines_names)
    should_continue_with_exclusion = request_confirmation()

    if should_continue_with_exclusion:
        pipelines_service.delete_pipelines_by_ids(pipelines_ids)
        return

    print('No pipelines deleted, exiting...')

def get_prop_to_filter_by():
    print('Which property should the pipelines be filtered by? default is "name"')
    return request_user_input('name')

def get_prop_value_to_filter_by():
    print('Which value should be used to filter this property? default is "disabled"')
    return request_user_input('disabled')

def request_user_input(default):
    user_input = input()
    if not user_input:
        return default
    return user_input

def display_retrieved_pipelines(disabled_pipelines_names):
    print(f'The following pipelines will be excluded: ')
    for pipeline in disabled_pipelines_names:
        print(pipeline)

def request_confirmation():
    print('Proceed with pipelines exclusion? [y/n]')
    user_input = input()
    return user_input == 'y'

if __name__ == "__main__":
    main()


import sys

import torch
from karies.config import ModelTypes
from karies.models import U_Net_Model

sys.path.append("/cluster/group/karies_2022/Simone/karies/karies-models/")
from configs.UNet.baseline_unet_config import model_config as config_unet

# ##### HYPERPARAMS #####
batch_size = 4

ENCODER_WEIGHTS_FOLDER = "/cluster/group/karies_2022/Simone/karies/karies-models/BYOL/unet_pixpro_weights/"
ENCODER_WEIGHTS_FILE = "encoder_batch_size_50_crop_size_56_epoch_400.pth"

config_unet["name"] = "PixPro_" + ENCODER_WEIGHTS_FILE
config_unet["fix_random_seed"] = 42
config_unet["unet_type"] = ModelTypes.U_Net
config_unet["batch_size"] = batch_size

# ##### TRAINING #####
unet = U_Net_Model(config_unet, load=False)
ENCODER_WEIGHTS = ENCODER_WEIGHTS_FOLDER + ENCODER_WEIGHTS_FILE
unet.model.encoder.load_state_dict(torch.load(ENCODER_WEIGHTS))

unet.train()


#! /usr/bin/env python
# -*- coding:utf8 -*-

import pandas as pd

from common.persistence.dbutils import create_session
from common.persistence.dao import IndexConstituentConfigDao
from framework.pagination import Pagination
from framework.localconcurrent import ProcessExcecutor

def update_index_constituent_config():
    """
    把20230309-finley 小于10的错误临时更新为正常
    Returns
    -------

    """
    session = create_session()
    results = session.execute("select distinct tscode, date from stock_validation_result where validation_code = '20230309-finley' and result = 1 and issue_count <= 10 order by tscode, date")
    update_list = list(map(lambda item: (item[0], item[1]), results))
    pagination = Pagination(update_list, page_size=50)
    while pagination.has_next():
        date_list = pagination.next()
        ProcessExcecutor(10).execute(update, date_list)

def update(*args):
    stock = args[0][0]
    date = args[0][1]
    date = date[0:4] + '-' + date[4:6] + '-' + date[6:8]
    print('Handle stock:{0} and date:{1}'.format(stock, date))
    constituent_config_dao = IndexConstituentConfigDao()
    constituent_config_dao.update_status(date, stock, 0)

if __name__ == '__main__':
    # update_index_constituent_config()

    # data = pd.read_hdf('D:\\liuli\\data\\original\\stock\\stock_data_tick_2016-01-14.h5')
    # print(data.shape)
    # print(data.columns)

    data = pd.read_hdf('D:\\liuli\\data\\original\\stock\\stock_data_1d.h5')
    print(data.shape)
    print(data.columns)
    print(data)
    print(data.loc['2005-01-04'])


import pymongo
from pymongo import MongoClient
import os

import discord
from discord.ext import commands

import random
import asyncio
import itertools
import sys
import traceback
from async_timeout import timeout
from functools import partial
from collections import deque
from youtube_dl import YoutubeDL
import youtube_dl
import datetime

cluster = MongoClient(os.environ['mongoURL'])
db = cluster["qwiooData"]
collection = db["userData"]

# Suppress noise about console usage from errors
youtube_dl.utils.bug_reports_message = lambda: ''

class VoiceError(Exception):
    pass

client = commands.Bot(command_prefix='??',
                      intents=discord.Intents.all(),
                      activity=discord.Game(name="music | ??help"))
ytdlopts = {
    'format': 'bestaudio/best',
    'extractaudio': True,
    'audioformat': 'mp3',
    'outtmpl': 'downloads/%(extractor)s-%(id)s-%(title)s.%(ext)s',
    'restrictfilenames': True,
    'noplaylist': True,
    'nocheckcertificate': True,
    'ignoreerrors': False,
    'logtostderr': False,
    'quiet': True,
    'no_warnings': True,
    'default_search': 'auto',
    'source_address': '0.0.0.0'  # ipv6 addresses cause issues sometimes
}

ffmpegopts = {
    'before_options': '-reconnect 1 -reconnect_streamed 1 -reconnect_delay_max 5',
    'options': '-vn'
}

ytdl = YoutubeDL(ytdlopts)

class YTDLError(Exception):
  pass

class VoiceConnectionError(commands.CommandError):
    """Custom Exception class for connection errors."""

class InvalidVoiceChannel(VoiceConnectionError):
    """Exception for cases of invalid Voice Channels."""

class YTDLSource(discord.PCMVolumeTransformer):

    def __init__(self, source, *, data, requester):
        super().__init__(source)
        self.requester = requester

        self.title = data.get('title')
        self.web_url = data.get('webpage_url')
        self.duration = self.parse_duration(int(data.get('duration')))
        self.thumbnail = data.get('thumbnail')

    def __getitem__(self, item: str):
        return self.__getattribute__(item)

    @classmethod
    async def create_source(cls, ctx, search: str, *, loop=None, download=False):
        loop = loop or asyncio.get_event_loop()

        to_run = partial(ytdl.extract_info, url=search, download=download)
        data = await loop.run_in_executor(None, to_run)

        if 'entries' in data:
            # take first item from a playlist
            data = data['entries'][0]

        embed = discord.Embed(title="", description=f"Queued [{data['title']}]({data['webpage_url']}) {datetime.timedelta(seconds=(data['duration']))} [{ctx.author.mention}]", color=0xF8C514)

        embed.set_thumbnail(url=data['thumbnail'])
        await ctx.send(embed=embed)

        if download:
            source = ytdl.prepare_filename(data)
        else:
            return {'webpage_url': data['webpage_url'], 'requester': ctx.author, 'title': data['title'], 'duration': str(datetime.timedelta(seconds=(data['duration'])))}

        return cls(discord.FFmpegPCMAudio(source), data=data, requester=ctx.author)

    @classmethod
    async def regather_stream(cls, data, *, loop):

        loop = loop or asyncio.get_event_loop()
        requester = data['requester']

        to_run = partial(ytdl.extract_info, url=data['webpage_url'], download=False)
        data = await loop.run_in_executor(None, to_run)

        return cls(discord.FFmpegPCMAudio(data['url']), data=data, requester=requester)

    @staticmethod
    def parse_duration(duration: int):
        minutes, seconds = divmod(duration, 60)
        hours, minutes = divmod(minutes, 60)
        days, hours = divmod(hours, 24)

        duration = []
        if days > 0:
            duration.append('{} days'.format(days))
        if hours > 0:
            duration.append('{} hr'.format(hours))
        if minutes > 0:
            duration.append('{} min'.format(minutes))
        if seconds > 0:
            duration.append('{} sec'.format(seconds))

        return ', '.join(duration)

class Spotify(commands.Cog):
    def getTrackID(self, track):
        track = sp.track(track)
        return track["id"]
    def getPlaylistTrackIDs(self, playlist_id):
        ids = []
        playlist = sp.playlist(playlist_id)
        for item in playlist['tracks']['items']:
            track = item['track']
            ids.append(track['id'])
        return ids

class MusicPlayer(commands.Cog):
    __slots__ = ('bot', '_guild', '_channel', 'voice_client', '_cog', 'queue', 'next', 'current', 'np', 'volume')

    def __init__(self, ctx):
        self.bot = ctx.bot
        self._guild = ctx.guild
        self._channel = ctx.channel

        self.voice_client = ctx.voice_client

        self._cog = ctx.cog

        self.queue = asyncio.Queue()
        self.next = asyncio.Event()
        self._loop = False
        self.np = None  # Now playing message
        self.volume = .69
        self.current = None

        ctx.bot.loop.create_task(self.player_loop())

    @property
    def loop(self):
        return self._loop

    @loop.setter
    def loop(self, value: bool):
        self._loop = value

    async def player_loop(self):

        """Our main player loop."""
        await self.bot.wait_until_ready()

        while not self.bot.is_closed():
            self.next.clear()
            if not self._loop:
              try:
                    # Wait for the next song. If we timeout - cancel the player and disconnect
                    async with timeout(300):
                        source = await self.queue.get()
              except asyncio.TimeoutError:
                    return self.destroy(self._guild)

            if self.loop:
                  self.queue.join(source)

            if not isinstance(source, YTDLSource):
                # regather to prevent stream expiration

                try:
                    source = await YTDLSource.regather_stream(source, loop=self.bot.loop)
                except Exception as e:
                    await self._channel.send(f'There was an error processing your song.\n'
                                             f'```css\n[{e}]\n```')
                    continue

            source.volume = self.volume
            self.current = source

            self._guild.voice_client.play(source, after=lambda _: self.bot.loop.call_soon_threadsafe(self.next.set))

            embed = discord.Embed(title="", description=f"[{source.title}]({source.web_url}) [{source.requester.mention}] \n`{source.duration}`", color=0xF8C514)

            embed.set_author(icon_url=self.bot.user.avatar_url, name=f"Now Playing 🎶")

            embed.set_thumbnail(url=source.thumbnail)
            await self._channel.send(embed=embed)

            await self.next.wait()
            source.cleanup()

            self.current = None

    def destroy(self, guild):
        """Disconnect and cleanup the player."""
        return self.bot.loop.create_task(self._cog.cleanup(guild))

class Music(commands.Cog):
    """Music related commands."""

    __slots__ = ('bot', 'players')

    def __init__(self, bot):
        self.bot = bot
        self.players = {}

    async def cleanup(self, guild):
        try:
            await guild.voice_client.disconnect()
        except AttributeError:
            pass

        try:
            del self.players[guild.id]
        except KeyError:
            pass

    async def __local_check(self, ctx):
        """A local check which applies to all commands in this cog."""
        if not ctx.guild:
            raise commands.NoPrivateMessage
        return True

    async def __error(self, ctx, error):
        """A local error handler for all errors arising from commands in this cog."""
        if isinstance(error, commands.NoPrivateMessage):
            try:
                return await ctx.send('This command can not be used in Private Messages.')
            except discord.HTTPException:
                pass
        elif isinstance(error, InvalidVoiceChannel):
            await ctx.send('Error connecting to Voice Channel. '
                           'Please make sure you are in a valid channel or provide me with one')

        print('Ignoring exception in command {}:'.format(ctx.command), file=sys.stderr)
        traceback.print_exception(type(error), error, error.__traceback__, file=sys.stderr)

    def get_player(self, ctx):
        """Retrieve the guild player, or generate one."""
        try:
            player = self.players[ctx.guild.id]
        except KeyError:
            player = MusicPlayer(ctx)
            self.players[ctx.guild.id] = player

        return player

    def _playlist(self, search: str):
      """Returns a dict with all Playlist entries"""
      ydl_opts = {
        'ignoreerrors': True,
        'quit': True,
        'source_address': '0.0.0.0'
      }

      with youtube_dl.YoutubeDL(ydl_opts) as ydl:
          playlist_dict = ydl.extract_info(search, download=False)

          playlistTitle = playlist_dict['title']

          playlist = dict()
          for video in playlist_dict['entries']:
              print()

              if not video:
                  print('ERROR: Unable to get info. Coninuing...')
                  continue

              for prop in ['id', 'title']:
                  print(prop, '--', video.get(prop))
                  playlist[video.get('title')] = 'https://www.youtube.com/watch?v=' + video.get('id')

          return playlist, playlistTitle

    @commands.command(name='join', aliases=['connect', 'j'], invoke_without_subcommand=True)
    async def connect_(self, ctx, *, channel: discord.VoiceChannel=None):
        if not channel:
            try:
                channel = ctx.author.voice.channel
            except AttributeError:
                embed = discord.Embed(title="", description="No channel to join. Please call `,join` from a voice channel.", color=0xF8C514)
                await ctx.send(embed=embed)
                raise InvalidVoiceChannel('No channel to join. Please either specify a valid channel or join one.')

        vc = ctx.voice_client

        if vc:
            if vc.channel.id == channel.id:
                return
            try:
                await vc.move_to(channel)
            except asyncio.TimeoutError:
                raise VoiceConnectionError(f'Moving to channel: <{channel}> timed out.')
        else:
            try:
                await channel.connect()
            except asyncio.TimeoutError:
                raise VoiceConnectionError(f'Connecting to channel: <{channel}> timed out.')
        if (random.randint(0, 1) == 0):
            await ctx.message.add_reaction('👍')
        await ctx.send(f'**Joined `{channel}`**')

    @commands.command(name='loop', aliases=['l'])
    async def _loop(self,ctx):
      player = self.get_player(ctx)
      vc = ctx.voice_client
      if not vc or not vc.is_playing:
            return await ctx.send('Nothing being played at the moment.')
      player.loop = not player.loop
      await ctx.message.add_reaction('✅')
      await ctx.send("Loop is now: " + str(player.loop))

    @commands.command(name='shuffle', aliases=['sf'])
    async def _shuffle(self, ctx):
        #player = self.get_player(ctx)
        player = self.get_player(ctx)

        if len(list(player.queue._queue)) > 1:
          player.queue._queue = random.sample(list(player.queue._queue), len(list(player.queue._queue)))
          player.queue._queue = deque(player.queue._queue)
          print(player.queue._queue)
          await ctx.send("Shuffled queue")
          await ctx.invoke(self.queue_info)

        else:
          await ctx.send("Empty Queue or queue size is 1.")

    @commands.command(name='play', aliases=['sing','p'])
    async def play_(self, ctx, *, search: str):

        playlist = 0
        await ctx.trigger_typing()
        player = self.get_player(ctx)

        vc = ctx.voice_client

        if not vc:
            await ctx.invoke(self.connect_)

        #check if link is a playlist
        # if search link has spotify in it

        if search.__contains__('&list='):
          playlist = 1
          await ctx.send('retrieving playlist...')
          async with ctx.typing():
            playlist, playlistTitle = self._playlist(search)
            for _title, _link in playlist.items():
                try:
                    source = await YTDLSource.create_source(ctx, _link, loop=self.bot.loop)
                    await player.queue.put(source)

                except YTDLError as e:
                    await ctx.send('An error occurred while processing this request: {}'.format(str(e)))

            await ctx.send(f'Enqueued `{playlist.__len__()}` songs from **{playlistTitle}**')

        #continue as usual
        player = self.get_player(ctx)
        if playlist == 1:
          pass
        elif playlist == 0:
          source = await YTDLSource.create_source(ctx, search, loop=self.bot.loop, download=False)
          await player.queue.put(source)

    @commands.command(name='pause')
    async def pause_(self, ctx):
        vc = ctx.voice_client

        if not vc or not vc.is_playing():
            embed = discord.Embed(title="", description="I am currently not playing anything", color=0xF8C514)
            return await ctx.send(embed=embed)
        elif vc.is_paused():
            return

        vc.pause()
        await ctx.send("Paused ⏸")

    @commands.command(name='resume')
    async def resume_(self, ctx):
        vc = ctx.voice_client

        if not vc or not vc.is_connected():
            embed = discord.Embed(title="", description="I'm not connected to a voice channel", color=0xF8C514)
            return await ctx.send(embed=embed)
        elif not vc.is_paused():
            return

        vc.resume()
        await ctx.send("Resuming ⏯")

    @commands.command(name='skip', aliases=['next','s'])
    async def skip_(self, ctx):
        vc = ctx.voice_client
        if not vc or not vc.is_connected():
            embed = discord.Embed(title="", description="I'm not connected to a voice channel", color=0xF8C514)
            return await ctx.send(embed=embed)

        if vc.is_paused():
            pass
        elif not vc.is_playing():
            return

        vc.stop()

    #TO DO
    #better error handling
    #command to add specific favs to queue
    #command to clear favourites
    @commands.command(name = 'favourite', aliases = ['f', 'fav'])
    async def favourite_(self, ctx, *, search):

        user = ctx.author
        myquery = { "_id": user.id}

        ydl_opts = {
          'format': 'bestaudio', 'noplaylist':'True'
        }
        with youtube_dl.YoutubeDL(ydl_opts) as ydl:
          data = ydl.extract_info(f"ytsearch:{search}", download=False)['entries'][0]
        print(data)
        dict = {
            'title': data['title'],
            'webpage_url': data['webpage_url'],
            'duration': str(datetime.timedelta(seconds=(data['duration'])))
          }
        print(dict)
        try:
          if (collection.count_documents(myquery) == 0):
            post = {"_id": user.id, "favourite": [dict]}
            collection.insert_one(post)

          else:
            collection.update_one({'_id' :user.id}, {"$push" : {'favourite': dict}})

          await ctx.channel.send("Added "+ data['title'] +  " to favourites")
        except discord.ext.commands.errors.MissingRequiredArgument:
          await ctx.channel.send("Include song after command")

    @commands.command(name = 'favourites', aliases = ['fs', 'favs'])
    async def favourites_(self, ctx):
        user = ctx.author
        myquery = { "_id": user.id }
        x = collection.find_one(myquery)

        try:
          favs = x['favourite']
          fmt = '\n'.join(f"`{(favs.index(_))}.` [{_['title']}]({_['webpage_url']}) | ` {_['duration']} `\n" for _ in favs)

          embed=discord.Embed(title="Favourites", description=fmt, color=0xF8C514)
          embed.set_footer(text=f"{ctx.author.display_name}", icon_url=ctx.author.avatar_url)
          await ctx.send(embed=embed)

        except TypeError as e:
          await ctx.send(e)

    @commands.command(name= 'removefavourite', aliases = ['rf', 'removefav', 'remfav'])
    async def removeFavourite_(self,ctx, pos : int=None):
      """Removes specified song from user's favourites"""
      user = ctx.author
      myquery = { "_id": user.id}

      #check if user exists in database
      if (collection.count_documents(myquery) > 0):

        x = collection.find_one(myquery)
        favs = x['favourite']

        try:
          await ctx.channel.send("removed " + str(favs[pos]['title']))
          del favs[pos]
          collection.replace_one({'_id' :user.id}, {'favourite': favs})

        except IndexError as e:
          await ctx.channel.send(e)

      else:
        await ctx.channel.send("use ??f to favourite songs")

    @commands.command(name='playfavourites', aliases=['pf','playfav'])
    async def play_favourites_(self,ctx,*, range):
        user = ctx.author
        myquery = { "_id": user.id}

      #check if user exists in database
        if (collection.count_documents(myquery) > 0):
          try:
            x = collection.find_one(myquery)
            favs = x['favourite']

            if range == 'all' or range == 'a':
              to_play = favs
            elif ',' in range:
              start, stop = range.split(',')

              to_play = favs[int(start) : (int(stop)+1)]

            elif len(str(range)) >= 1:

              to_play = favs[(int(range)) : (int(range)+1)]
              print(to_play)

            for song in to_play:
              print(song)
              link = song['webpage_url']
              print(type(link))

              await ctx.invoke(self.bot.get_command('play'), search = link)

          except ValueError as e:
            await ctx.channel.send(e)

        else:
          await ctx.channel.send("use ??f to favourite songs")
#--------------------------------------------------------------------------------

    @commands.command(name='remove', aliases=['rm', 'rem', 'r'])
    async def remove_(self, ctx, pos : int=None):
        """Removes specified song from queue"""

        vc = ctx.voice_client

        if not vc or not vc.is_connected():
            embed = discord.Embed(title="", description="I'm not connected to a voice channel", color=0xF8C514)
            return await ctx.send(embed=embed)

        player = self.get_player(ctx)
        if pos == None:
            player.queue._queue.pop()
        else:
            try:
                s = player.queue._queue[pos-1]
                del player.queue._queue[pos-1]
                embed = discord.Embed(title="", description=f"Removed [{s['title']}]({s['webpage_url']}) [{s['requester'].mention}]", color=0xF8C514)
                await ctx.send(embed=embed)
            except:
                embed = discord.Embed(title="", description=f'Could not find a track for "{pos}"', color=0xF8C514)
                await ctx.send(embed=embed)

    @commands.command(name='clear', aliases=['clr', 'cl', 'cr'])
    async def clear_(self, ctx):
        """Deletes entire queue of upcoming songs."""

        vc = ctx.voice_client

        if not vc or not vc.is_connected():
            embed = discord.Embed(title="", description="I'm not connected to a voice channel", color=0xF8C514)
            return await ctx.send(embed=embed)

        player = self.get_player(ctx)
        player.queue._queue.clear()
        await ctx.send('💣 **Cleared**')

    @commands.command(name='queue', aliases=['q', 'playlist', 'que'])
    async def queue_info(self, ctx):
        """Retrieve a basic queue of upcoming songs."""
        vc = ctx.voice_client

        if not vc or not vc.is_connected():
            embed = discord.Embed(title="", description="I'm not connected to a voice channel", color=0xF8C514)
            return await ctx.send(embed=embed)

        player = self.get_player(ctx)
        if player.queue.empty():
            embed = discord.Embed(title="", description="queue is empty", color=0xF8C514)
            return await ctx.send(embed=embed)

        print(player.queue._queue)
        # Grabs the songs in the queue...
        upcoming = list(itertools.islice(player.queue._queue, 0, int(len(player.queue._queue))))

        fmt = '\n'.join(f"`{(upcoming.index(_)) + 1}.` [{_['title']}]({_['webpage_url']}) | ` {_['duration']} Requested by: {_['requester']}`\n" for _ in upcoming)

        fmt = f"\n__Now Playing__:\n[{vc.source.title}]({vc.source.web_url}) | ` {vc.source.duration} Requested by: {vc.source.requester}`\n\n__Up Next:__\n" + fmt + f"\n**{len(upcoming)} songs in queue**"

        embed = discord.Embed(title=f'Queue for {ctx.guild.name}', description=fmt, color=0xF8C514)
        embed.set_footer(text=f"{ctx.author.display_name}", icon_url=ctx.author.avatar_url)

        await ctx.send(embed=embed)

    @commands.command(name='np', aliases=['song', 'current', 'currentsong', 'playing'],invoke_without_command=True)
    async def now_playing_(self, ctx):
        """Display information about the currently playing song."""
        vc = ctx.voice_client

        if not vc or not vc.is_connected():
            embed = discord.Embed(title="", description="I'm not connected to a voice channel", color=0xF8C514)
            return await ctx.send(embed=embed)

        player = self.get_player(ctx)
        if not player.current:
            embed = discord.Embed(title="", description="I am currently not playing anything", color=0xF8C514)
            return await ctx.send(embed=embed)

        embed = discord.Embed(title="", description=f"[{vc.source.title}]({vc.source.web_url}) [{vc.source.requester.mention}]\n`{vc.source.duration}`", color=0xF8C514)
        embed.set_author(icon_url=self.bot.user.avatar_url, name=f"Now Playing 🎶")
        embed.set_image(url=vc.source.thumbnail)
        msg = await ctx.send(embed=embed)

        first_run = True
        while True:
            if first_run:
              reactmoji = ["⏸","▶","⏩","🎶",'🔊']
              for react in reactmoji:
                          await msg.add_reaction(react)

              def check_react(reaction, user):
                      if reaction.message.id != msg.id:
                          return False
                      if user != ctx.message.author:
                          return False
                      if str(reaction.emoji) not in reactmoji:
                          return False
                      return True
              try:
                      res, user = await self.bot.wait_for('reaction_add', check=check_react)
              except asyncio.TimeoutError:
                      return await msg.clear_reactions()

              if '⏸' in str(res.emoji):
                      print('<<paused>>')
                      ctx.voice_client.pause()

              elif '▶' in str(res.emoji):
                      print('<<resumed>>')
                      ctx.voice_client.resume()

              elif '⏩' in str(res.emoji):
                      print('<<skipped>>')
                      if not vc or not vc.is_connected():
                          embed = discord.Embed(title="", description="I'm not connected to a voice channel", color=0xF8C514)
                          return await ctx.send(embed=embed)

                      if vc.is_paused():
                          pass
                      elif not vc.is_playing():
                          return

                      vc.stop()

              elif '🎶' in str(res.emoji):
                      print('<<Displaying queue>>')
                      await ctx.invoke(self.queue_info)

              elif '🔊' in str(res.emoji):
                      print('<<Displaying volume>>')
                      await ctx.invoke(self.change_volume)

    @commands.command(name='volume', aliases=['vol', 'v'])
    async def change_volume(self, ctx, *, vol: float=None):
        """Change the player volume.
        Parameters
        ------------
        volume: float or int [Required]
            The volume to set the player to in percentage. This must be between 1 and 100.
        """
        vc = ctx.voice_client

        if not vc or not vc.is_connected():
            embed = discord.Embed(title="", description="I am not currently connected to voice", color=0xF8C514)
            return await ctx.send(embed=embed)

        if not vol:
            embed = discord.Embed(title="", description=f"🔊 **{(vc.source.volume)*100}%**", color=0xF8C514)
            return await ctx.send(embed=embed)

        if not 0 < vol < 101:
            embed = discord.Embed(title="", description="Please enter a value between 1 and 100", color=0xF8C514)
            return await ctx.send(embed=embed)

        player = self.get_player(ctx)

        if vc.source:
            vc.source.volume = vol / 100

        player.volume = vol / 100
        embed = discord.Embed(title="", description=f'**`{ctx.author}`** set the volume to **{vol}%**', color=0xF8C514)
        await ctx.send(embed=embed)

    @commands.command(name='leave', aliases=["stop", "dc", "disconnect", "bye"])
    async def leave_(self, ctx):
        """Stop the currently playing song and destroy the player.
        !Warning!
            This will destroy the player assigned to your guild, also deleting any queued songs and settings.
        """
        vc = ctx.voice_client

        if not vc or not vc.is_connected():
            embed = discord.Embed(title="", description="I'm not connected to a voice channel", color=0xF8C514)
            return await ctx.send(embed=embed)

        if (random.randint(0, 1) == 0):
            await ctx.message.add_reaction('👋')
        await ctx.send('**Successfully disconnected**')

        await self.cleanup(ctx.guild)

def setup(bot):
    bot.add_cog(Music(bot))


import requests
from bs4 import BeautifulSoup
import re
headers = {
    'User-Agent':'Mozilla/4.0(compatible;MSIE 5.5;Windows NT)',}
url = 'http://www.biqukan.com/3_3039/1351331.html'
html = requests.get(url,headers=headers)
# 获取到网页的文本格式
html = html.text
soup = BeautifulSoup(html,'lxml')
# 选择器得到的是一个列表，只有一个元素，要选取第一个
zhangjie = soup.select('div.content h1')[0]
# 得到章节的文本内容
zhangjie = zhangjie.text
print(zhangjie)
# 查找属性是id为'content'且class为'showtxt'的标签
text = soup.find(id='content', class_='showtxt')
# 得到文本内容
text=text.text
# 将得到的文本存入文件中
with open('xiaoshuo.txt','w',encoding='utf-8') as f:
    f.write(zhangjie)
    f.write(text)


from typing import Iterable
import numpy as np
import pandas as pd

# functions checks if there are different values in a series. If thats the case it returns NaN, else it returns the unique value.
def series_single_unique_val_or_nan(series: pd.Series):
    arr = series.unique()
    length = len(arr)
    return arr[0] if length == 1 else np.NaN # if only one unique elem given return it, else return NaN

# Returns a dictionary where for every distinct element in the iterable a mapping to a unique letter is given. Complexity is linear.
def get_to_char_map(iterable: Iterable) -> dict[str, str]:
    res = {}
    letter = 'a'
    for x in iterable:
        if not x in res:
            res[x] = letter
            letter = chr(ord(letter)+1)
    return res


from bs4 import BeautifulSoup
# from PIL import Image
from selenium import webdriver
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.common.by import By
from webdriver_manager.chrome import ChromeDriverManager
# import os
import urllib.request as urllib2

def getsize(uri):
    file = urllib2.urlopen(uri)
    return len(file.read())

s = Service(ChromeDriverManager().install())
driver = webdriver.Chrome(service = s)
driver.get("https://www.instagram.com/lauweded/")

f = open("results.html", "w+")
f.write(str(driver.find_element(By.CLASS_NAME, 'SCxLW').get_attribute('innerHTML').encode('utf-8').strip()))
f.close()

with open("results.html") as fp :
    soup = BeautifulSoup(fp, "lxml")

# Get numbers of followers

stats_account = soup.find_all(attrs={'class': 'k9GMp'})
followers = 0
publications = 0

for index, num in enumerate(stats_account) :
    # Numbers of publications
    if index == 0 :
        publications = num.find('span').text
    # Numbers of followers
    if index == 1 :
        followers = num.find('span').text

print(followers)

# Get images and their sizes

images = soup.find_all('img')
sum_sizes = 0
image_count = 0

for image in images:
    # print(image.get('src'))
    size = getsize(image.get('src')) * 4
    size_str = str(size) + ' bytes'
    # print(size_str)
    sum_sizes += size
    image_count += 1

# print(sum_sizes)
# print(image_count)

# Calculations

# poids d'une image en moyenne
average_size = float(sum_sizes) / float(image_count)
print(average_size)

# Nombre de publications par la taille moyenne
size_tot = float(publications) * average_size
print(size_tot)

# poids total de toutes les images base sur la taille moyenne en fonction du nombre de followers
size_tot += float(followers) * average_size
print(size_tot)

# Consommation electrique ; (poids total / 1E+06) * 1E-03
elect = (size_tot / 1000000.0) * 0.001
print(elect)

# Consommation en CO2
emission = elect * 0.276
print(emission)