from __future__ import absolute_import, division, print_function
from simtbx.diffBragg.stage_two_utils import PAR_from_params
import os
import sys
from libtbx.mpi4py import MPI
COMM = MPI.COMM_WORLD
from dials.array_family import flex
import numpy as np
try:
    import pandas
except ImportError:
    print("Pandas is required. Install using 'libtbx.python -m pip install pandas'")
    exit()
from xfel.merging.application.utils.memory_usage import get_memory_usage
from simtbx.diffBragg.refiners.stage_two_refiner import StageTwoRefiner
from simtbx.diffBragg import utils
from simtbx.diffBragg import hopper_utils
from dxtbx.model.experiment_list import ExperimentListFactory
from simtbx.diffBragg.prep_stage2_input import prep_dataframe
from cctbx import miller, crystal
import logging

LOGGER = logging.getLogger("diffBragg.main")

def global_refiner_from_parameters(params):
    launcher = RefineLauncher(params)
    # TODO read on each rank, or read and broadcast ?
    LOGGER.info("EVENT: read input pickle")
    pandas_table = pandas.read_pickle(params.pandas_table)
    LOGGER.info("EVENT: BEGIN prep dataframe")
    if params.prep_time > 0:
        work_distribution = prep_dataframe(pandas_table)
    LOGGER.info("EVENT: DONE prep dataframe")
    return launcher.launch_refiner(pandas_table, work_distribution=work_distribution)

class RefineLauncher:

    def __init__(self, params):
        self.params = self.check_parameter_integrity(params)
        self.n_shots_on_rank = None
        self.df = None
        self.Modelers = {}
        self.Hi = {}
        self.Hi_asu = {}
        self.symbol = None
        self.DEVICE_ID = 0

    @property
    def NPIX_TO_ALLOC(self):
        return self._NPIX_TO_ALLOC

    @NPIX_TO_ALLOC.setter
    def NPIX_TO_ALLOC(self, val):
        assert val> 0 or val == -1
        self._NPIX_TO_ALLOC = int(val)

    @staticmethod
    def check_parameter_integrity(params):
        if params.refiner.max_calls is None or len(params.refiner.max_calls) == 0:
            raise ValueError("Cannot refine because params.refiner.max_calls is empty")

        if os.environ.get("DIFFBRAGG_CUDA") is not None:
            params.refiner.use_cuda = True

        return params

    @property
    def num_shots_on_rank(self):
        return len(self.Modelers)

    def _init_panel_group_information(self, detector):
        # default is one group per panel:
        self.panel_group_from_id = {pid: 0 for pid in range(len(detector))}
        # if specified in a file, then overwrite:
        if self.params.refiner.panel_group_file is not None:
            self.panel_group_from_id = utils.load_panel_group_file(self.params.refiner.panel_group_file)
            if not self.panel_group_from_id:
                raise ValueError("Loading panel group file %s  produced an panel group dict!"
                                 % self.params.refiner.panel_group_file)

        # how many unique panel groups :
        panel_groups = set(self.panel_group_from_id.values())
        self.n_panel_groups = len(panel_groups)

        # make dict where the key is the panel group id, and the value is a list of the panel ids within the group
        panels_per_group = {group_id: [] for group_id in panel_groups}
        for pid in self.panel_group_from_id:
            group_id = self.panel_group_from_id[pid]
            panels_per_group[group_id].append(pid)

        # we should rotate each panel in a group about the same reference point
        # Make a dict where key is panel id, and the  value is a reference origin
        self.panel_reference_from_id = {}
        for pid in self.panel_group_from_id:
            group_id = self.panel_group_from_id[pid]

            # take as reference, the origin of the first panel in the group
            reference_panel = detector[panels_per_group[group_id][0]]
            self.panel_reference_from_id[pid] = reference_panel.get_origin()

    def _init_simulator(self, expt, miller_data):
        self.SIM = utils.simulator_for_refinement(expt, self.params)
        # note self.SIM.D is a now diffBragg instance
        # update the miller data ?
        if miller_data is not None:
            self.SIM.crystal.miller_array = miller_data.as_amplitude_array()
            self.SIM.update_Fhkl_tuple()

    @staticmethod
    def _check_experiment_integrity(expt):
        for model in ["crystal", "detector", "beam", "imageset"]:
            if not hasattr(expt, model):
                raise ValueError("No %s in experiment, exiting. " % model)

    def launch_refiner(self, pandas_table, miller_data=None, work_distribution=None):
        self.load_inputs(pandas_table, miller_data=miller_data, work_distribution=work_distribution)
        LOGGER.info("EVENT: launch refiner")
        self._launch()
        return self.RUC

    def load_inputs(self, pandas_table, miller_data=None, work_distribution=None, refls_key='predictions'):
        """

        :param pandas_table: contains path to the experiments (pandas column exp_name) to be loaded
            the pandas table is expected to have been written by diffBragg.hopper or
            diffBragg.hopper_process . See method save_to_pandas in simtbx/command_line/hopper.py
            For example, if the outputdir of diffBragg.hopper was set to `all_shots`, then
            there should be a golder all_shots/pandas created which contains all of the per-shot pandas
            dataframes. They should be concatenated as follows, forming a suitable argument for this method
            >> import glob,pandas
            >> fnames = glob.glob("all_shots/pandas/rank*/*pkl")
            >> df = pandas.concat([ pandas.read_pickle(f) for f in fnames])
            >> df.reset_index(inplace=True, drop=True)
            >> df.to_pickle("all_shots.pkl")
            >> # Then later, as part of an MPI application, the following will load all data:
            >> RefineLauncher_instance.load_inputs(df, refls_key="stage1_refls")

        :param miller_data: Optional miller array for the structure factor component of the model
        :param refls_key: key specifying the reflection tables in the pandas table
            Modeled pixels will lie in shoeboxes centered on each x,y,z in xyzobs.px.value
        :return:
        """
        COMM.Barrier()
        num_exp = len(pandas_table)
        first_exper_file = pandas_table.exp_name.values[0]
        detector = ExperimentListFactory.from_json_file(first_exper_file, check_format=False)[0].detector
        if detector is None and self.params.refiner.reference_geom is None:
            raise RuntimeError("No detector in experiment, must provide a reference geom.")
        # TODO verify all shots have the same detector ?
        if self.params.refiner.reference_geom is not None:
            detector = ExperimentListFactory.from_json_file(self.params.refiner.reference_geom, check_format=False)[0].detector
            print("Using reference geom from expt %s" % self.params.refiner.reference_geom)

        if COMM.size > num_exp:
            raise ValueError("Requested %d MPI ranks to process %d shots. Reduce number of ranks to %d"
                             % (COMM.size, num_exp, num_exp))
        self._init_panel_group_information(detector)

        self.verbose = False
        if COMM.rank == 0:
            self.verbose = self.params.refiner.verbose > 0
            if self.params.refiner.gather_dir is not None and not os.path.exists(self.params.refiner.gather_dir):
              os.makedirs(self.params.refiner.gather_dir)
              LOGGER.info("MADE GATHER DIR %s" % self.params.refiner.gather_dir)
        COMM.barrier()
        shot_idx = 0  # each rank keeps index of the shots local to it
        rank_panel_groups_refined = set()
        exper_names = pandas_table.exp_name
        assert len(exper_names) == len(set(exper_names))
        # TODO assert all exper are single-file, probably way before this point
        if work_distribution is None:
            worklist = range(COMM.rank, len(exper_names), COMM.size)
        else:
            worklist = work_distribution[COMM.rank]
        LOGGER.info("EVENT: begin loading inputs")
        for i_exp in worklist:
            exper_name = exper_names[i_exp]
            LOGGER.info("EVENT: BEGIN loading experiment list")
            expt_list = ExperimentListFactory.from_json_file(exper_name, check_format=self.params.refiner.check_expt_format)
            LOGGER.info("EVENT: DONE loading experiment list")
            if len(expt_list) != 1:
                print("Input experiments need to have length 1, %s does not" % exper_name)
            expt = expt_list[0]
            expt.detector = detector  # in case of supplied ref geom
            self._check_experiment_integrity(expt)

            exper_dataframe = pandas_table.query("exp_name=='%s'" % exper_name)

            refl_name = exper_dataframe[refls_key].values[0]
            refls = flex.reflection_table.from_file(refl_name)
            # FIXME need to remove (0,0,0) bboxes

            try:
                good_sel = flex.bool([h != (0, 0, 0) for h in list(refls["miller_index"])])
                refls = refls.select(good_sel)
            except KeyError:
                pass

            #UcellMan = utils.manager_from_crystal(expt.crystal)
            opt_uc_param = exper_dataframe[["a","b","c","al","be","ga"]].values[0]
            UcellMan = utils.manager_from_params(opt_uc_param)

            if self.symbol is None:
                if self.params.refiner.force_symbol is not None:
                    self.symbol = self.params.refiner.force_symbol
                else:
                    self.symbol = expt.crystal.get_space_group().type().lookup_symbol()
                LOGGER.info("Set space group symbol: %s" % self.symbol)
            else:
                if self.params.refiner.force_symbol is None:
                    if expt.crystal.get_space_group().type().lookup_symbol() != self.symbol:
                        raise ValueError("Crystals should all have the same space group symmetry")

            if shot_idx == 0:  # each rank initializes a simulator only once
                if self.params.simulator.init_scale != 1:
                    print("WARNING: For stage_two , it is assumed that total scale is stored in the pandas dataframe")
                    print("WARNING: resetting params.simulator.init_scale to 1!")
                    self.params.simulator.init_scale = 1
                self._init_simulator(expt, miller_data)
                if self.params.profile:
                    self.SIM.record_timings = True
                if self.params.refiner.stage_two.Fref_mtzname is not None:
                    self.Fref = utils.open_mtz(self.params.refiner.stage_two.Fref_mtzname,
                                               self.params.refiner.stage_two.Fref_mtzcol)

            if "miller_index" in list(refls.keys()):
                is_allowed = flex.bool(len(refls), True)
                allowed_hkls = set(self.SIM.crystal.miller_array.indices())
                uc = self.SIM.crystal.dxtbx_crystal.get_unit_cell().parameters()
                symb = self.SIM.crystal.space_group_info.type().lookup_symbol()
                sym = crystal.symmetry(uc, symb)
                mset = miller.set(sym, refls['miller_index'],True)
                op = mset.change_of_basis_op_to_primitive_setting()
                mset_p = mset.change_basis(op)
                refl_hkls_p1 = mset_p.indices()

                for i_ref in range(len(refls)):
                    hkl_p1 = refl_hkls_p1[i_ref]
                    if hkl_p1 not in allowed_hkls:
                        is_allowed[i_ref] = False
                refls = refls.select(is_allowed)

            LOGGER.info("EVENT: LOADING ROI DATA")
            shot_modeler = hopper_utils.DataModeler(self.params)
            shot_modeler.exper_name = exper_name
            shot_modeler.refl_name = refl_name
            shot_modeler.rank = COMM.rank
            if self.params.refiner.load_data_from_refl:
                gathered = shot_modeler.GatherFromReflectionTable(expt, refls, sg_symbol=self.symbol)
            else:
                gathered = shot_modeler.GatherFromExperiment(expt, refls, sg_symbol=self.symbol)
            if not gathered:
                raise IOError("Failed to gather data from experiment %s", exper_name)
                COMM.abort()

            if self.params.refiner.gather_dir is not None:
                gathered_name = os.path.splitext(os.path.basename(exper_name))[0]
                gathered_name += "_withData.refl"
                gathered_name = os.path.join(self.params.refiner.gather_dir, gathered_name )
                shot_modeler.dump_gathered_to_refl(gathered_name, do_xyobs_sanity_check=False) #True)
                LOGGER.info("SAVED ROI DATA TO %s" % gathered_name)
                if self.params.refiner.test_gathered_file:
                    all_data = shot_modeler.all_data.copy()
                    all_roi_id = shot_modeler.roi_id.copy()
                    all_bg = shot_modeler.all_background.copy()
                    all_trusted = shot_modeler.all_trusted.copy()
                    all_pids = np.array(shot_modeler.pids)
                    all_rois = np.array(shot_modeler.rois)
                    new_Modeler = hopper_utils.DataModeler(self.params)
                    assert new_Modeler.GatherFromReflectionTable(exper_name, gathered_name, sg_symbol=self.symbol)
                    assert np.allclose(new_Modeler.all_data, all_data)
                    assert np.allclose(new_Modeler.all_background, all_bg)
                    assert np.allclose(new_Modeler.rois, all_rois)
                    assert np.allclose(new_Modeler.pids, all_pids)
                    assert np.allclose(new_Modeler.all_trusted, all_trusted)
                    assert np.allclose(new_Modeler.roi_id, all_roi_id)
                    LOGGER.info("Gathered file approved!")

            self.Hi[shot_idx] = shot_modeler.Hi
            self.Hi_asu[shot_idx] = shot_modeler.Hi_asu

            LOGGER.info("EVENT: DONE LOADING ROI")
            shot_modeler.ucell_man = UcellMan
            self.SIM.num_ucell_param = len(shot_modeler.ucell_man.variables)  # for convenience

            loaded_spectra = False
            if self.params.spectrum_from_imageset:
                try:
                    shot_spectra = hopper_utils.downsamp_spec(self.SIM, self.params, expt, return_and_dont_set=True)
                    loaded_spectra = True
                except Exception as err:
                    LOGGER.warning("spectrum_from_imageset is set to True, however failed to load spectra: %s" % err)
                    loaded_spectra = False

            if not loaded_spectra:
                if "spectrum_filename" in list(exper_dataframe) and exper_dataframe.spectrum_filename.values[0] is not None:
                    shot_spectra = utils.load_spectra_from_dataframe(exper_dataframe)
                    LOGGER.debug("Loaded specta from %s" % exper_dataframe.spectrum_filename.values[0])
                    shot_modeler.spec_name = exper_dataframe.spectrum_filename.values[0]

                else:
                    total_flux = exper_dataframe.total_flux.values[0]
                    if total_flux is None:
                        total_flux = self.params.simulator.total_flux
                    shot_spectra = [(expt.beam.get_wavelength(), total_flux)]

            shot_modeler.spectra = shot_spectra
            if self.params.refiner.gather_dir is not None and not self.params.refiner.load_data_from_refl:
                spec_wave, spec_weights = map(np.array, zip(*shot_spectra))
                spec_filename = os.path.splitext(os.path.basename(exper_name))[0]
                spec_filename = os.path.join(self.params.refiner.gather_dir, spec_filename+".lam")
                utils.save_spectra_file(spec_filename, spec_wave, spec_weights)
                LOGGER.info("saved spectra filename %s" % spec_filename)

            LOGGER.info("Will simulate %d energy channels" % len(shot_spectra))

            if "detz_shift_mm" in list(exper_dataframe):
                shot_modeler.originZ_init = exper_dataframe.detz_shift_mm.values[0]*1e-3
            else:
                shot_modeler.originZ_init = 0
            shot_modeler.exper_name = exper_name
            shot_modeler.refl_name = refl_name

            shot_panel_groups_refined = self.determine_refined_panel_groups(shot_modeler.pids)
            rank_panel_groups_refined = rank_panel_groups_refined.union(set(shot_panel_groups_refined))

            shot_idx += 1
            if COMM.rank == 0:
                self._mem_usage()
                print("Finished loading image %d / %d" % (i_exp+1, len(exper_names)), flush=True)

            shot_modeler.PAR = PAR_from_params(self.params, expt, best=exper_dataframe)
            self.Modelers[i_exp] = shot_modeler

        LOGGER.info("DONE LOADING DATA; ENTER BARRIER")
        COMM.Barrier()
        LOGGER.info("DONE LOADING DATA; EXIT BARRIER")
        self.shot_roi_darkRMS = None

        # TODO warn that per_spot_scale refinement not intended for ensemble mode
        all_refined_groups = COMM.gather(rank_panel_groups_refined)
        panel_groups_refined = None
        if COMM.rank == 0:
            panel_groups_refined = set()
            for set_of_panels in all_refined_groups:
                panel_groups_refined = panel_groups_refined.union(set_of_panels)
        self.panel_groups_refined = list(COMM.bcast(panel_groups_refined))

        LOGGER.info("EVENT: Gathering global HKL information")
        try:
            self._gather_Hi_information()
        except TypeError:
            pass
        LOGGER.info("EVENT: FINISHED gather global HKL information")
        if self.params.roi.cache_dir_only:
            print("Done creating cache directory and cache_dir_only=True, so goodbye.")
            sys.exit()

        # in case of GPU
        LOGGER.info("BEGIN DETERMINE MAX PIX")
        self.NPIX_TO_ALLOC = self._determine_per_rank_max_num_pix()
        # TODO in case of randomize devices, shouldnt this be total max across all ranks?
        n = COMM.gather(self.NPIX_TO_ALLOC)
        if COMM.rank == 0:
            n = max(n)
        self.NPIX_TO_ALLOC = COMM.bcast(n)
        LOGGER.info("DONE DETERMINE MAX PIX")

        self.DEVICE_ID = COMM.rank % self.params.refiner.num_devices
        self._mem_usage()

    def _mem_usage(self):
        memMB = get_memory_usage()
        import socket
        host = socket.gethostname()
        print("Rank 0 reporting memory usage: %f GB on Rank 0 node %s" % (memMB / 1e3, host))

    def determine_refined_panel_groups(self, pids):
        refined_groups = []
        #assert len(pids) == len(selection_flags)
        for i, pid in enumerate(pids):
            if self.panel_group_from_id[pid] not in refined_groups:
                refined_groups.append(self.panel_group_from_id[pid])
        return refined_groups

    def _determine_per_rank_max_num_pix(self):
        max_npix = 0
        for i_shot in self.Modelers:
            modeler = self.Modelers[i_shot]
            x1, x2, y1, y2 = map(np.array, zip(*modeler.rois))
            npix = np.sum((x2-x1)*(y2-y1))
            max_npix = max(npix, max_npix)
            print("Rank %d, shot %d has %d pixels" % (COMM.rank, i_shot+1, npix))
        print("Rank %d, max pix to be modeled: %d" % (COMM.rank, max_npix))
        return max_npix

    def _try_loading_spectrum_filelist(self):
        file_list = None
        fpath = self.params.simulator.spectrum.filename_list
        if fpath is not None:
            file_list = [l.strip() for l in open(fpath, "r").readlines()]
            assert all([len(l.split()) == 1 for l in file_list]), "weird spectrum file %s"% fpath
        return file_list

    def _gather_Hi_information(self):
        nshots_on_this_rank = len(self.Hi)
        # aggregate all miller indices
        self.Hi_all_ranks, self.Hi_asu_all_ranks = [], []
        # TODO assert list types are stored in Hi and Hi_asu
        for i_shot in self.Hi: #range(nshots_on_this_rank):
            self.Hi_all_ranks += self.Hi[i_shot]
            self.Hi_asu_all_ranks += self.Hi_asu[i_shot]
        self.Hi_all_ranks = COMM.reduce(self.Hi_all_ranks)
        self.Hi_all_ranks = COMM.bcast(self.Hi_all_ranks)

        self.Hi_asu_all_ranks = COMM.reduce(self.Hi_asu_all_ranks)
        self.Hi_asu_all_ranks = COMM.bcast(self.Hi_asu_all_ranks)

        marr_unique_h = self._get_unique_Hi()

        # this will map the measured miller indices to their index in the LBFGS parameter array self.x
        self.idx_from_asu = {h: i for i, h in enumerate(set(self.Hi_asu_all_ranks))}
        # we will need the inverse map during refinement to update the miller array in diffBragg, so we cache it here
        self.asu_from_idx = {i: h for i, h in enumerate(set(self.Hi_asu_all_ranks))}

        self.num_hkl_global = len(self.idx_from_asu)

        fres = marr_unique_h.d_spacings()
        self.res_from_asu = {h: res for h, res in zip(fres.indices(), fres.data())}

    def _get_unique_Hi(self):
        COMM.barrier()
        if COMM.rank == 0:
            from cctbx.crystal import symmetry
            from cctbx import miller
            from cctbx.array_family import flex as cctbx_flex

            ii = list(self.Modelers.keys())[0]
            uc = self.Modelers[ii].ucell_man
            params = uc.a, uc.b, uc.c, uc.al * 180 / np.pi, uc.be * 180 / np.pi, uc.ga * 180 / np.pi
            if self.params.refiner.force_unit_cell is not None:
                params = self.params.refiner.force_unit_cell
            symm = symmetry(unit_cell=params, space_group_symbol=self.symbol)
            hi_asu_flex = cctbx_flex.miller_index(self.Hi_asu_all_ranks)
            mset = miller.set(symm, hi_asu_flex, anomalous_flag=True)
            marr = miller.array(mset)
            binner = marr.setup_binner(d_max=self.params.refiner.stage_two.d_max, d_min=self.params.refiner.stage_two.d_min,
                                       n_bins=self.params.refiner.stage_two.n_bin)
            from collections import Counter
            print("Average multiplicities:")
            print("<><><><><><><><><><><><>")
            for i_bin in range(self.params.refiner.stage_two.n_bin - 1):
                dmax, dmin = binner.bin_d_range(i_bin + 1)
                F_in_bin = marr.resolution_filter(d_max=dmax, d_min=dmin)
                multi_in_bin = np.array(list(Counter(F_in_bin.indices()).values()))
                print("%2.5g-%2.5g : Multiplicity=%.4f" % (dmax, dmin, multi_in_bin.mean()))
                for ii in range(1, 100, 8):
                    print("\t %d refls with multi %d" % (sum(multi_in_bin == ii), ii))

            print("Overall completeness\n<><><><><><><><>")
            symm = symmetry(unit_cell=params, space_group_symbol=self.symbol)
            hi_flex_unique = cctbx_flex.miller_index(list(set(self.Hi_asu_all_ranks)))
            mset = miller.set(symm, hi_flex_unique, anomalous_flag=True)
            self.binner = mset.setup_binner(d_min=self.params.refiner.stage_two.d_min,
                                            d_max=self.params.refiner.stage_two.d_max,
                                            n_bins=self.params.refiner.stage_two.n_bin)
            mset.completeness(use_binning=True).show()
            marr_unique_h = miller.array(mset)
            print("Rank %d: total miller vars=%d" % (COMM.rank, len(set(self.Hi_asu_all_ranks))))
        else:
            marr_unique_h = None

        marr_unique_h = COMM.bcast(marr_unique_h)
        return marr_unique_h

    def _launch(self):
        """
        Usually this method should be modified when new features are added to refinement
        """
        # TODO return None or refiner instance
        LOGGER.info("begin _launch")
        x_init = None
        nmacro = self.params.refiner.num_macro_cycles
        n_trials = len(self.params.refiner.max_calls)
        for i_trial in range(n_trials*nmacro):

            self.RUC = StageTwoRefiner(self.Modelers, self.symbol, self.params)

            if self.will_refine(self.params.refiner.refine_spot_scale):
                self.RUC.refine_crystal_scale = (self.params.refiner.refine_spot_scale*nmacro)[i_trial]

            if self.will_refine(self.params.refiner.refine_Fcell):
                self.RUC.refine_Fcell = (self.params.refiner.refine_Fcell*nmacro)[i_trial]

            self.RUC.panel_group_from_id = self.panel_group_from_id
            self.RUC.panel_reference_from_id = self.panel_reference_from_id
            self.RUC.panel_groups_being_refined = self.panel_groups_refined

            # TODO verify not refining Fcell in case of local refiner
            self.RUC.max_calls = (self.params.refiner.max_calls*nmacro)[i_trial]
            self.RUC.x_init = x_init
            self.RUC.ignore_line_search_failed_step_at_lower_bound = True  # TODO: why was this necessary?

            # plot things
            self.RUC.trial_id = i_trial

            self.RUC.log_fcells = True
            self.RUC.request_diag_once = False
            self.RUC.trad_conv = True
            self.RUC.idx_from_asu = self.idx_from_asu
            self.RUC.asu_from_idx = self.asu_from_idx

            self.RUC.S = self.SIM
            self.RUC.restart_file = self.params.refiner.io.restart_file
            self.RUC.S.update_nanoBragg_instance('update_oversample_during_refinement',
                                                 self.params.refiner.update_oversample_during_refinement)
            self.RUC.S.update_nanoBragg_instance("Npix_to_allocate", self.NPIX_TO_ALLOC)
            self.RUC.S.update_nanoBragg_instance('device_Id', self.DEVICE_ID)
            self.RUC.use_curvatures_threshold = self.params.refiner.use_curvatures_threshold
            if not self.params.refiner.curvatures:
                self.RUC.S.update_nanoBragg_instance('compute_curvatures', False)
            if COMM.rank==0:
                self.RUC.S.update_nanoBragg_instance('verbose', self.params.refiner.verbose)

            LOGGER.info("_launch run setup")
            self.RUC.run(setup_only=True)
            LOGGER.info("_launch done run setup")
            # for debug purposes:
            #if not self.params.refiner.quiet:
            #    print("\n<><><><><><><><>TRIAL %d refinement status:" % i_trial)
            #    self.RUC.S.D.print_if_refining()

            self.RUC.num_positive_curvatures = 0
            self.RUC.use_curvatures = self.params.refiner.start_with_curvatures
            self.RUC.hit_break_to_use_curvatures = False

            # selection flags set here:
            #self.RUC.selection_flags = self.shot_selection_flags
            #if self.params.refiner.res_ranges is not None:
            #    assert self.shot_reso is not None, "cant set reso flags is rlp is not in refl tables"
            #    nshots = len(self.shot_selection_flags)
            #    more_sel_flags = {}
            #    res_ranges = utils.parse_reso_string(self.params.refiner.res_ranges)
            #    for i_shot in range(nshots):
            #        rhigh, rlow = (res_ranges*nmacro)[i_trial]
            #        sel_flags = self.shot_selection_flags[i_shot]
            #        res_flags = [rhigh < r < rlow for r in self.shot_reso[i_shot]]
            #        more_sel_flags[i_shot] = [flag1 and flag2 for flag1,flag2 in zip(sel_flags, res_flags)]
            #    self.RUC.selection_flags = more_sel_flags

            LOGGER.info("_launcher runno setup")
            self.RUC.run(setup=False)
            LOGGER.info("_launcher done runno setup")
            if self.RUC.hit_break_to_use_curvatures:
                self.RUC.fix_params_with_negative_curvature = False
                self.RUC.num_positive_curvatures = 0
                self.RUC.use_curvatures = True
                self.RUC.run(setup=False)

            if self.RUC.hit_break_signal:
                if self.params.profile:
                    self.RUC.S.D.show_timings(self.RUC.rank)
                self.RUC._MPI_barrier()
                break

            if self.params.refiner.debug_pixel_panelfastslow is not None:
                utils.show_diffBragg_state(self.RUC.S.D, self.params.refiner.debug_pixel_panelfastslow)
                s = self.RUC._get_spot_scale(0)
                print("refiner spot scale=%f" % (s**2))

            x_init = self.RUC.x

            if self.params.profile:
                self.RUC.S.D.show_timings(self.RUC.rank)
            if os.environ.get("DIFFBRAGG_USE_CUDA") is not None:
                self.RUC.S.D.gpu_free()

    def will_refine(self, param):
        return param is not None and any(param)


import torch
from torch import nn
from torch import optim
from torch.nn import functional as F
import numpy as np
import logging
import os
import re
import json
from collections import defaultdict
from copy import deepcopy
from pprint import pformat
import random

def pad(seqs, emb, device, pad=0):
    lens = [len(s) for s in seqs]
    max_len = max(lens)
    padded = torch.LongTensor([s + (max_len-l) * [pad] for s, l in zip(seqs, lens)])
    return emb(padded.to(device)), lens

def run_rnn(rnn, inputs, lens):
    # sort by lens
    order = np.argsort(lens)[::-1].tolist()
    reindexed = inputs.index_select(0, inputs.data.new(order).long())
    reindexed_lens = [lens[i] for i in order]
    packed = nn.utils.rnn.pack_padded_sequence(reindexed, reindexed_lens, batch_first=True)
    outputs, _ = rnn(packed)
    padded, _ = nn.utils.rnn.pad_packed_sequence(outputs, batch_first=True, padding_value=0.)
    reverse_order = np.argsort(order).tolist()
    recovered = padded.index_select(0, inputs.data.new(reverse_order).long())
    # reindexed_lens = [lens[i] for i in order]
    # recovered_lens = [reindexed_lens[i] for i in reverse_order]
    # assert recovered_lens == lens
    return recovered

def attend(seq, cond, lens):
    """
    attend over the sequences `seq` using the condition `cond`.
    """
    scores = cond.unsqueeze(1).expand_as(seq).mul(seq).sum(2)
    max_len = max(lens)
    for i, l in enumerate(lens):
        if l < max_len:
            scores.data[i, l:] = -np.inf
    scores = F.softmax(scores, dim=1)
    context = scores.unsqueeze(2).expand_as(seq).mul(seq).sum(1)
    return context, scores

class FixedEmbedding(nn.Embedding):
    """
    this is the same as `nn.Embedding` but detaches the result from the graph and has dropout after lookup.
    """

    def __init__(self, *args, dropout=0, **kwargs):
        super().__init__(*args, **kwargs)
        self.dropout = dropout

    def forward(self, *args, **kwargs):
        out = super().forward(*args, **kwargs)
        out.detach_()
        return F.dropout(out, self.dropout, self.training)

class SelfAttention(nn.Module):
    """
    scores each element of the sequence with a linear layer and uses the normalized scores to compute a context over the sequence.
    """

    def __init__(self, d_hid, dropout=0.):
        super().__init__()
        self.scorer = nn.Linear(d_hid, 1)
        self.dropout = nn.Dropout(dropout)

    def forward(self, inp, lens):
        batch_size, seq_len, d_feat = inp.size()
        inp = self.dropout(inp)
        scores = self.scorer(inp.contiguous().view(-1, d_feat)).view(batch_size, seq_len)
        max_len = max(lens)
        for i, l in enumerate(lens):
            if l < max_len:
                scores.data[i, l:] = -np.inf
        scores = F.softmax(scores, dim=1)
        context = scores.unsqueeze(2).expand_as(inp).mul(inp).sum(1)
        return context

class GLADEncoder(nn.Module):
    """
    the GLAD encoder described in https://arxiv.org/abs/1805.09655.
    """

    def __init__(self, din, dhid, slots, dropout=None):
        super().__init__()
        self.dropout = dropout or {}
        self.global_rnn = nn.LSTM(din, dhid, bidirectional=True, batch_first=True)
        self.global_selfattn = SelfAttention(2 * dhid, dropout=self.dropout.get('selfattn', 0.))
        for s in slots:
            setattr(self, '{}_rnn'.format(s), nn.LSTM(din, dhid, bidirectional=True, batch_first=True, dropout=self.dropout.get('rnn', 0.)))
            setattr(self, '{}_selfattn'.format(s), SelfAttention(2*dhid, dropout=self.dropout.get('selfattn', 0.)))
        self.slots = slots
        self.beta_raw = nn.Parameter(torch.Tensor(len(slots)))
        nn.init.uniform_(self.beta_raw, -0.01, 0.01)

    def beta(self, slot):
        return F.sigmoid(self.beta_raw[self.slots.index(slot)])

    def forward(self, x, x_len, slot, default_dropout=0.2):
        local_rnn = getattr(self, '{}_rnn'.format(slot))
        local_selfattn = getattr(self, '{}_selfattn'.format(slot))
        beta = self.beta(slot)
        local_h = run_rnn(local_rnn, x, x_len)
        global_h = run_rnn(self.global_rnn, x, x_len)
        h = F.dropout(local_h, self.dropout.get('local', default_dropout), self.training) * beta + F.dropout(global_h, self.dropout.get('global', default_dropout), self.training) * (1-beta)
        c = F.dropout(local_selfattn(h, x_len), self.dropout.get('local', default_dropout), self.training) * beta + F.dropout(self.global_selfattn(h, x_len), self.dropout.get('global', default_dropout), self.training) * (1-beta)
        return h, c

from transformers import BertTokenizer
tokenizer_path = "bert/"
lower_case = False
cache_path = "data/cache"

class Model(nn.Module):
    """
    the GLAD model described in https://arxiv.org/abs/1805.09655.
    """

    def __init__(self, args, ontology, vocab, des):
        super().__init__()
        self.optimizer = None
        self.args = args
        self.vocab = vocab
        self.ontology = ontology
        self.emb_fixed = FixedEmbedding(len(vocab), args.demb, dropout=args.dropout.get('emb', 0.2))

        self.utt_encoder = GLADEncoder(args.demb, args.dhid, self.ontology.slots, dropout=args.dropout)
        self.act_encoder = GLADEncoder(args.demb, args.dhid, self.ontology.slots, dropout=args.dropout)
        self.ont_encoder = GLADEncoder(args.demb, args.dhid, self.ontology.slots, dropout=args.dropout)
        self.utt_scorer = nn.Linear(2 * args.dhid, 1)
        self.score_weight = nn.Parameter(torch.Tensor([0.5]))
        self.new_weight = nn.Parameter(torch.Tensor([0.5]))

        self.news_scorer = nn.Linear(2*args.dhid, 1)

        # self.tokenizer =BertTokenizer.from_pretrained(tokenizer_path, do_lower_case=lower_case, cache_dir=cache_path)
        drop_parameter = 0.1
        mask_types = 2
        self.dropout = nn.Dropout(drop_parameter)
        self.mask_classifier = nn.Linear(args.demb, mask_types)

    @property
    def device(self):
        if self.args.gpu is not None and torch.cuda.is_available():
            return torch.device('cuda')
        else:
            return torch.device('cpu')

    def set_optimizer(self):
        self.optimizer = optim.Adam(self.parameters(), lr=self.args.lr)

    def load_emb(self, Eword):
        new = self.emb_fixed.weight.data.new
        self.emb_fixed.weight.data.copy_(new(Eword))

    def mask_pred(self, num_transcript, utterance, utterance_len, mask, emb):
        batch_size = len(mask)
        # print(mask)
        mask_clf_list = []
        for bs in range(batch_size):
            # curr_nt = torch.LongTensor(num_transcript[bs])
            curr_len = utterance_len[bs]
            curr_utt = utterance[bs][:curr_len,:]
            curr_mask = torch.LongTensor(mask[bs]).to(self.device).unsqueeze(0)      #[1, 7, 6]
            # emb_nt = emb(curr_nt.to(self.device))       #[6, 400]

            # 传入了被GladEncoder编码的utterance 根据len进行切割 去掉padding的部分
            m = (curr_mask.unsqueeze(-1) == 0 ).float() * (-1e30)  #[1, 7, 6, 1] 7种mask的可能性 6的长度
            t = curr_utt.unsqueeze(0).repeat(1, curr_mask.shape[1], 1, 1)
            entity_span_pool = m + t  # 1 7 6 400
            entity_span_pool = entity_span_pool.max(dim=2)[0]  #[1, 7, 400]
            entity_repr = self.dropout(entity_span_pool)
            mask_clf = self.mask_classifier(entity_repr.squeeze(0))
            mask_clf = torch.softmax(mask_clf, dim=1)
            mask_clf_list.append(mask_clf)
            # print(entity_span_pool.size())
        return mask_clf_list

    def mask_pred_multi(self, num_transcript, utterance, utterance_len, mask, emb):
        batch_size = len(mask)
        # print(mask)
        mask_clf_list = []
        for bs in range(batch_size):
            # curr_nt = torch.LongTensor(num_transcript[bs])
            curr_len = utterance_len[bs]
            curr_utt = utterance[bs][:curr_len,:]   #   [6, 400]
            curr_mask = torch.LongTensor(mask[bs]).to(self.device)      #[7, 6]

            # 传入了被GladEncoder编码的utterance 根据len进行切割 去掉padding的部分
            curr_mask = curr_mask.float()
            m = torch.matmul(curr_mask, curr_utt)
            entity_repr = self.dropout(m)
            mask_clf = self.mask_classifier(entity_repr)
            mask_clf = torch.softmax(mask_clf, dim=1)
            mask_clf_list.append(mask_clf)
        return mask_clf_list

    def gen_true_mask(self, mask, pos):
        # mask是整体mask pos和neg都有     而pos只是pos
        batch_size = len(mask)
        mask_label = []
        for bs in range(batch_size):
            curr_mask = mask[bs]
            curr_pos = pos[bs]
            m_label = [[False, True] for i in range(len(mask[bs]))]  #第一列mask为True 第二列mask为False
            for p in curr_pos:
                true_index = curr_mask.index(p)
                m_label[true_index][0] = True
                m_label[true_index][1] = False
            m_label = torch.LongTensor(m_label).to(self.device).float()
            mask_label.append(m_label)
            # print(m_label)  #如果是未打乱之前的 那么应该只有前几个为True 后面全是False
        return mask_label

    def get_pred_mask(self, utterance, mask_clf, entity_mask, threshold = 0.4):
        entity_mask_tensor_list = []
        for em in entity_mask:
            entity_mask_tensor_list.append(torch.LongTensor(em).to(self.device))
        weight = []
        for mc, em in zip(mask_clf, entity_mask_tensor_list):
            pos_clf = mc[:,0]>threshold
            weighted_mask_list = []
            for mask_weight, mask in zip(pos_clf, em):
                if mask_weight:
                    weighted_mask_list.append(mask)
            if len(weighted_mask_list)==0:
                return None
            weighted_mask_tensor = torch.stack(weighted_mask_list)
            weighted_mask_tensor = torch.sum(weighted_mask_tensor, dim=0)
            weight.append(weighted_mask_tensor)

        # pad
        max_weight_len = max([w.size(0) for w in weight])
        for weight_index in range(len(weight)):
            pad_num = max_weight_len - weight[weight_index].size(0)
            weight[weight_index] = torch.nn.functional.pad(weight[weight_index], pad=(0, pad_num), mode='constant',
                                                           value=0)
        weight_tensor = torch.stack(weight)
        weighted_utterance = self.make_weighted(weight_tensor, utterance)

        return weighted_utterance

    # 相乘mask纵列相加 横行归一化
    def cal_weighted_mask(self, mask_clf, entity_mask_list):
        entity_mask_tensor_list = []
        weight = []
        for em in entity_mask_list:
            entity_mask_tensor_list.append(torch.LongTensor(em).to(self.device))

        for mc, emtl in zip(mask_clf, entity_mask_tensor_list):
            pos_clf = mc[:,0]
            weighted_mask_list = []
            for mask_weight, mask in zip(pos_clf, emtl):
                weighted_mask_list.append(mask_weight*mask)
            weighted_mask_tensor = torch.stack(weighted_mask_list)
            weighted_mask_sum = torch.sum(weighted_mask_tensor, dim=0) #tensor([0.0000, 0.9350, 1.4798, 1.7114, 1.0551, 0.0000]
            weighted_mask_sum = F.softmax(weighted_mask_sum)  #TODO 是否有必要softmax需要思考

            # 到这一步 weight_mask_sum就可以表明predicted出的一个句子中每个token关于是否有slot这件事的概率值 也可以理解为信息量
            weight.append(weighted_mask_sum)

        # pad
        max_weight_len = max([w.size(0) for w in weight])
        for weight_index in range(len(weight)):
            pad_num = max_weight_len-weight[weight_index].size(0)
            weight[weight_index] = torch.nn.functional.pad(weight[weight_index], pad=(0,pad_num), mode='constant', value=0)
        weight_tensor = torch.stack(weight)

        return weight_tensor

    def make_weighted(self, mask_weight, thing):
        weighted_t = []
        for mw, t in zip(mask_weight, thing):  # thing 可替换成替他等长度的东西
            t_mw = t * mw.unsqueeze(0).T
            weighted_t.append(t_mw)
        weighted_thing = torch.stack(weighted_t)
        return weighted_thing

    def concat_mask(self, pos, neg):
        if not pos:
            return neg
        elif not neg:
            return pos
        else:
            return pos+neg

    def forward(self, batch):
        # convert to variables and look up embeddings
        eos = self.vocab.word2index('<eos>')
        utterance, utterance_len = pad([e.num['transcript'] for e in batch], self.emb_fixed, self.device, pad=eos)
        # num_transcript =[e.num['transcript'] for e in batch]

        entity_mask, num_transcript, pos_mask = [], [], []
        for e in batch:
            # print(e.to_dict()['mask'])
            n_transcript = e.num['transcript']
            e_mask = self.concat_mask(e.to_dict()['mask']['pos'], e.to_dict()['mask']['neg'])  #还需要标注出pos
            p_mask =  e.to_dict()['mask']['pos']
            num_transcript.append(n_transcript)
            entity_mask.append(e_mask)
            pos_mask.append(p_mask)
        # entity mask是否需要shuffle 如果shuffle就在这统一shuffle mmp 不能shuffle 因为后续需要utterance和mask对应 TODO 除非弄个class 把这两个集成为一个对象中的两个属性
        # random.shuffle(entity_mask)
        # entity_mask_tensor = torch.LongTensor(entity_mask)  # cannot convert to tensor because every length of mask varies
        # print([e.to_dict()['mask'] for e in batch])

        acts = [pad(e.num['system_acts'], self.emb_fixed, self.device, pad=eos) for e in batch]
        # slot-value pairs under consideration
        ontology = {s: pad(v, self.emb_fixed, self.device, pad=eos) for s, v in self.ontology.num.items()}
        ys, exist = {}, {} # print(ontology['area'])  #'area', 'price range', 'request'
        # mask_clf_list[s] = self.mask_pred_multi(num_transcript, utterance, utterance_len, entity_mask, self.emb_fixed)
        mask_clf_list = self.mask_pred(num_transcript, utterance, utterance_len, entity_mask, self.emb_fixed)
        true_mask_list = self.gen_true_mask(entity_mask, pos_mask)
        # mask_weight = self.cal_weighted_mask(mask_clf_list, entity_mask)
        # weight_utterance = self.make_weighted(mask_weight, utterance)

        # weighted_utt = self.get_pred_mask(utterance, mask_clf_list, entity_mask)
        # if weighted_utt is not None:
        #     utterance = torch.cat([utterance, weighted_utt], dim=1)   #expand 7,25,400 ->7,50,400

        for s in self.ontology.slots:
            # for each slot, compute the scores for each value
            H_utt, c_utt = self.utt_encoder(utterance, utterance_len, slot=s)  #Hutt 7 25 400 cutt7 400
            H_acts, C_acts = list(zip(*[self.act_encoder(a, a_len, slot=s) for a, a_len in acts]))  #Hacts 7*[1 3 400]
            _, C_vals = self.ont_encoder(ontology[s][0], ontology[s][1], slot=s)  #0 represetn seqs and 1 represent lens,  only encode the values. the slot is help to discriminate

            # think add position information here  # input: e.num['transcript']  utterance e['mask'](pos and neg) utterance_len
            # H_utt = self.make_weighted(mask_weight, H_utt) #给处理之后的Hutt施加以weight

            # compute the utterance score
            y_utts, q_utts, q_acts, q_news = [], [], [], []
            for c_val in C_vals:
                cond = c_val.unsqueeze(0).expand(len(batch), *c_val.size())
                q_utt, _ = attend(H_utt, cond, lens=utterance_len)  # H_utt 7,25,400    cond 7,400
                q_utts.append(q_utt)
            # print(torch.stack(q_utts, dim=1).size()) #7 2 400
            y_utts = self.utt_scorer(torch.stack(q_utts, dim=1)).squeeze(2)

            # compute the previous action score
            for i, C_act in enumerate(C_acts):
                seq = C_act.unsqueeze(0)
                cond = c_utt[i].unsqueeze(0)
                q_act, _ = attend(seq, cond, lens=[C_act.size(0)]) #act,seq 1,1,400     utt,cond 1,400
                q_acts.append(q_act)
            y_acts = torch.cat(q_acts, dim=0).mm(C_vals.transpose(0, 1))

            # val自注意力
            # for c_val in C_vals:
            #     cond = c_val.unsqueeze(0).expand(len(batch), *c_val.size())
            #     seq = cond.unsqueeze(1)
            #     q_new, _ = attend(seq, cond, lens=[seq.size(0)])
            #     q_news.append(q_new)
            # y_news = self.news_scorer(torch.stack(q_news, dim=1)).squeeze(2)

            # exist[s] = self.exist_scorer(torch.stack(q_utts, dim=1)).squeeze(2)
            # exist[s] = torch.softmax(exist[s], dim=1)
            # exist[s] = torch.index_select(exist[s], 1, torch.tensor([0]).to(self.device))

            output = y_utts + self.score_weight * y_acts
            # output = self.make_weighted(exist[s] * 2, output).squeeze(1)
            ys[s] = F.sigmoid(output)

        if self.training:
            # print(self.ontology.values) {'area': ['south', 'west'], 'price range': ['cheap', 'expensive'], 'request': ['address', 'phone']}
            # create label variable and compute loss
            labels = {s: [len(self.ontology.values[s]) * [0] for i in range(len(batch))] for s in self.ontology.slots}
            for i, e in enumerate(batch):
                for s, v in e.turn_label:
                    # print("%s %s %s"%(s, v, self.ontology.values[s].index(v)))  #area south 0
                    labels[s][i][self.ontology.values[s].index(v)] = 1  #牛啊 核心在s上 因为一句话虽然可能有三个turn_label 但是每种只有一个值 这样就做到了最终与utt数对齐
            labels = {s: torch.Tensor(m).to(self.device) for s, m in labels.items()}
            # exist_labels = {s: (torch.sum(labels[s],dim=1)!=0).float().unsqueeze(1) for s in labels}

            loss, mask_loss, exist_loss = 0, 0, 0
            for s in self.ontology.slots:
                loss += F.binary_cross_entropy(ys[s], labels[s])
                # exist_loss += F.binary_cross_entropy(exist[s], exist_labels[s])
            for index in range(len(true_mask_list)):
                m_loss = F.binary_cross_entropy(mask_clf_list[index], true_mask_list[index])
                mask_loss += m_loss

        else:
            loss = torch.Tensor([0]).to(self.device)
            mask_loss = torch.Tensor([0]).to(self.device)
            # exist_loss = torch.Tensor([0]).to(self.device)
        score = {s: v.data.tolist() for s, v in ys.items()}
        # mask_score = {s: v.data.tolist() for s,v in mask_clf_list.items()}
        # print(mask_score)

        return loss, score, mask_clf_list, true_mask_list

    def get_train_logger(self):
        logger = logging.getLogger('train-{}'.format(self.__class__.__name__))
        formatter = logging.Formatter('%(asctime)s [%(threadName)-12.12s] [%(levelname)-5.5s]  %(message)s')
        file_handler = logging.FileHandler(os.path.join(self.args.dout, 'train.log'))
        file_handler.setFormatter(formatter)
        logger.addHandler(file_handler)
        return logger

    def run_train(self, train, dev, args):
        # print(train.to_dict())
        track = defaultdict(list)
        iteration = 0
        best = {}
        logger = self.get_train_logger()
        if self.optimizer is None:
            self.set_optimizer()

        for epoch in range(args.epoch):
            logger.info('starting epoch {}'.format(epoch))

            # train and update parameters
            self.train()
            # print("train:")  是一个batch  但是在一个batch中含有多个turn和两个dialogue
            for batch in train.batch(batch_size=args.batch_size, shuffle=True, type="train"):
                # for t in batch:
                #     print(t.to_dict())
                iteration += 1   # the batch number
                self.zero_grad()
                loss, scores, mask_score, true_mask = self.forward(batch)
                loss.backward()
                self.optimizer.step()
                track['loss'].append(loss.item())

            # evalute on train and dev
            summary = {'iteration': iteration, 'epoch': epoch}
            #   k:loss.  multi-batch. divide by the batch num
            for k, v in track.items():
                summary[k] = sum(v) / len(v)
            print("================")
            summary.update({'eval_train_{}'.format(k): v for k, v in self.run_eval(train, args).items()})
            print("----------------")
            summary.update({'eval_dev_{}'.format(k): v for k, v in self.run_eval(dev, args).items()})
            print("================")

            print("")
            # do early stopping saves
            stop_key = 'eval_dev_{}'.format(args.stop)
            train_key = 'eval_train_{}'.format(args.stop)
            if best.get(stop_key, 0) <= summary[stop_key]:
                # print(best.get(stop_key, 0))
                # print(summary[stop_key])
                best_dev = '{:f}'.format(summary[stop_key])
                best_train = '{:f}'.format(summary[train_key])
                best.update(summary)
                self.save(best, identifier='epoch={epoch},iter={iteration},train_{key}={train},dev_{key}={dev}'.format(epoch=epoch, iteration=iteration, train=best_train, dev=best_dev, key=args.stop, ) )
                self.prune_saves()  # detele some model that unneccesary to save
                dev.record_preds(
                    preds=self.run_pred(dev, self.args)[0],
                    to_file=os.path.join(self.args.dout, 'dev.pred.json'),
                )
            summary.update({'best_{}'.format(k): v for k, v in best.items()})

            logger.info(pformat(summary))  #INFO:train-Model
            track.clear()

    def extract_predictions(self, scores, mask_score, threshold=0.4):
        batch_size = len(list(scores.values())[0])
        predictions = [set() for i in range(batch_size)]
        for s in self.ontology.slots:
            for i, p in enumerate(scores[s]):
                triggered = [(s, v, p_v) for v, p_v in zip(self.ontology.values[s], p) if p_v > 0.3]
                if s == 'request':
                    # we can have multiple requests predictions
                    predictions[i] |= set([(s, v) for s, v, p_v in triggered])
                elif triggered:
                    # only extract the top inform prediction
                    sort = sorted(triggered, key=lambda tup: tup[-1], reverse=True)
                    predictions[i].add((sort[0][0], sort[0][1]))

        # mask是与slot种类无关的 需要结合entity_mask及mask_score判断我预测哪些mask为True 然后将预测为True的mask与pos比对
        pred_mask = [ ms>threshold for ms in mask_score]
        # print(predictions)

        return predictions, pred_mask

    # prediction and evaluate
    def run_pred(self, dev, args):
        self.eval()
        predictions = []
        mask_predictions = []
        gold_masks = []
        for batch in dev.batch(batch_size=args.batch_size, type="eval "):
            loss, scores, mask_score, gold_mask = self.forward(batch)
            pred, pred_mask = self.extract_predictions(scores, mask_score)
            predictions += pred
            mask_predictions += pred_mask
            gold_masks += gold_mask

        return predictions, mask_predictions, gold_masks

    def run_eval(self, dev, args):
        predictions, mask_predictions, gold_mask = self.run_pred(dev, args)
        return dev.evaluate_preds(predictions, mask_predictions, gold_mask)

    def save_config(self):
        fname = '{}/config.json'.format(self.args.dout)
        with open(fname, 'wt') as f:
            # logging.info('saving config to {}'.format(fname))
            json.dump(vars(self.args), f, indent=2)

    @classmethod
    def load_config(cls, fname, ontology, **kwargs):
        with open(fname) as f:
            logging.info('loading config from {}'.format(fname))
            args = object()
            for k, v in json.load(f):
                setattr(args, k, kwargs.get(k, v))
        return cls(args, ontology)

    def save(self, summary, identifier):
        fname = '{}/{}.t7'.format(self.args.dout, identifier)
        logging.info('saving model to {}'.format(fname))
        state = {
            'args': vars(self.args),
            'model': self.state_dict(),
            'summary': summary,
            'optimizer': self.optimizer.state_dict(),
        }
        torch.save(state, fname)

    def load(self, fname):
        logging.info('loading model from {}'.format(fname))
        state = torch.load(fname)
        self.load_state_dict(state['model'])
        self.set_optimizer()
        self.optimizer.load_state_dict(state['optimizer'])

    def get_saves(self, directory=None):
        if directory is None:
            directory = self.args.dout
        files = [f for f in os.listdir(directory) if f.endswith('.t7')]
        scores = []
        for fname in files:
            re_str = r'dev_{}=([0-9\.]+)'.format(self.args.stop)
            dev_acc = re.findall(re_str, fname)
            if dev_acc:
                score = float(dev_acc[0].strip('.'))
                scores.append((score, os.path.join(directory, fname)))
        if not scores:
            raise Exception('No files found!')
        scores.sort(key=lambda tup: tup[0], reverse=True)  # what a nb! here sort the model for list according by score
        return scores

    def prune_saves(self, n_keep=5):
        scores_and_files = self.get_saves()
        # if the number of model in folder is larger than n_keep, this will delete
        if len(scores_and_files) > n_keep:
            for score, fname in scores_and_files[n_keep:]:
                try:
                    os.remove(fname)
                except:
                    print("delete save fail")

    def load_best_save(self, directory):
        if directory is None:
            directory = self.args.dout

        scores_and_files = self.get_saves(directory=directory)
        if scores_and_files:
            assert scores_and_files, 'no saves exist at {}'.format(directory)
            score, fname = scores_and_files[0]
            self.load(fname)


import struct
import socket
import threading
import time
from collections import deque
from my.QUIC.quic_client import QUICClient

# define format of HTTP 3 frame format
class Frame:
    def __init__(self):
        self.type: int = 0
        self.length: int = 0
        self.payload: bytes = b""

    @staticmethod
    def pack_frame(type: int, payload: bytes):
        length = len(payload)
        pack_str = f"!BI{length}s"
        return struct.pack(pack_str, type, length, payload)

    def unpack_frame(self, frame_header: bytes):
        self.type, self.length = struct.unpack("!BI", frame_header[:5])
        self.payload = frame_header[5:]

    def print_frame(self):
        print(f"type: {self.type}")
        print(f"length: {self.length}")
        print(f"payload: {self.payload}")

def parse_byte_header(header: bytes) -> dict:
    headers = {}
    header = header.decode()
    for line in header.split("\r\n"):
        if ":" in line:
            key, value = line.split(": ", 1)
            headers[key] = value
    return headers

class Request:
    def __init__(self, method, path, headers=None, body=None):
        self.method = "GET"
        self.path = path
        self.headers = headers if headers is not None else {}
        self.body = body

    def get_header_string(self):
        if self.headers is None:
            return ""
        header_string = ""
        for key, value in self.headers.items():
            header_string += f"{key}: {value}\r\n"
        return header_string

    def to_frame(self):
        payload = f"{self.method} {self.path} HTTP/2.0\r\n{self.get_header_string()}\r\n".encode()
        return Frame.pack_frame(1, payload)

class Response():
    def __init__(self, stream_id, headers = {}, status = "Not yet") -> None:
        self.stream_id = stream_id
        self.headers = headers

        self.status = status
        self.body = b""

        self.contents = deque()
        self.complete = False

    def get_headers(self):
        begin_time = time.time()
        while self.status == "Not yet":
            if time.time() - begin_time > 5:
                return None
        return self.headers

    def get_full_body(self): # used for handling short body
        begin_time = time.time()
        while not self.complete:
            if time.time() - begin_time > 5:
                return None
        if len(self.body) > 0:
            return self.body
        while len(self.contents) > 0:
            self.body += self.contents.popleft()
        return self.body # the full content of HTTP response body

    def get_stream_content(self): # used for handling long body
        begin_time = time.time()
        while len(self.contents) == 0: # contents is a buffer, busy waiting for new content
            if self.complete or time.time() - begin_time > 5: # if response is complete or timeout
                return None
            pass
        content = self.contents.popleft() # pop content from deque
        return content # the part content of the HTTP response body

class HTTPClient(): # For HTTP/2.X
    streams: dict[int, Response]
    def __init__(self) -> None:
        self.client = QUICClient()
        self.connected = False
        self.last_stream_id = -1
        self.streams = {}
        self.recv_thread = threading.Thread(target=self.recv_worker)
        self.recv_thread.start()

    @staticmethod
    def parse_url(url):
        if "http://" in url:
            url = url.replace("http://", "")
        parts = url.split("/")

        host = parts[0]

        # Split the host into hostname and port number
        host_parts = host.split(":")
        hostname = host_parts[0]  # The first part is the hostname

        if len(host_parts) > 1:
            port = int(host_parts[1])  # The second part is the port number
        else:
            port = 80  # Default port is 80 for HTTP

        # Join the remaining parts as the path
        path = "/" + "/".join(parts[1:])

        return hostname, port, path

    def form_request_frame(self, path, headers=None, body=None) -> bytes:
        headers = {
            "method": "GET",
            "path": path,
            "scheme": "http",
            "authority": "127.0.0.1:6789",
            "Content-Length": "0"
        }
        request = Request("GET", path, headers, body)

        return request.to_frame()

    def recv_worker(self):
        while True:
            recved = self.client.recv()
            frame = Frame()
            frame.unpack_frame(recved[1])

            if frame.length != len(frame.payload):
                print("length not match")
                continue

            stream_id = recved[0]
            is_end = recved[2]
            print(f"stream_id: {stream_id}, is_end: {is_end}")
            if stream_id not in self.streams:
                self.streams[stream_id] = Response(stream_id)
            if frame.type == 0x1: # headers frame
                self.streams[stream_id].headers = parse_byte_header(frame.payload)
            if frame.type == 0x0: # data frame
                self.streams[stream_id].contents.append(frame.payload)
                # print out length of deque with message
                # print(f"stream {stream_id} has {len(self.streams[stream_id].contents)} contents")
            if is_end == 0x1 and frame.type == 0x0: # end of stream and data frame
                self.streams[stream_id].complete = True
                self.streams[stream_id].status = "Complete"

    def get(self, url, headers=None):
        print("request url: ", url)
        host, port, path = self.parse_url(url)
        if not self.connected:
            self.client.connect((host, port))
            self.connected = True

        # Send the request and return the response
        self.last_stream_id += 2
        request_frame = self.form_request_frame(path, headers=headers)
        self.client.send(self.last_stream_id, request_frame, True)
        self.streams[self.last_stream_id] = Response(self.last_stream_id)

        return self.streams[self.last_stream_id]


from chalice.test import Client
from pytest import fixture
from app import app
import json

@fixture
def test_client():
    with Client(app) as client:
        yield client

def test_start_game(test_client):
    result = test_client.http.post(
        "/jogo/novo",
        headers={
            "Content-type": "application/json"
        },
        body=json.dumps({
            'nick_name': 'pepo'
        })
    )
    assert result.status_code == 201

def test_start_game_header_location(test_client):
    result = test_client.http.post(
        "/jogo/novo",
        headers={
            "Content-type": "application/json"
        },
        body=json.dumps({
            'nick_name': 'Alisson'
        })
    )
    headers = list(map(lambda x: x.lower(), result.headers))
    assert 'location' in headers

def test_start_game_get_game_id(test_client):
    result = test_client.http.post(
        "/jogo/novo",
        headers={
            "Content-type": "application/json"
        },
        body=json.dumps({
            'nick_name': 'Alisson'
        })
    )

    headers = {k.lower(): v for k, v in result.headers.items()}
    uriSplitted = headers['location'].split('/')

    assert len(uriSplitted) > 3


# -*- coding: utf-8 -*-

import xml.etree.ElementTree as et
import matplotlib.pyplot as plt
import os
import sys
import math
import matplotlib
# matplotlib.use('Agg')

def loadPos(xml, xpath, id):
    x = [[[], [], []], [[], [], []], [[], [], []]]
    y = [[[], [], []], [[], [], []], [[], [], []]]
    z = [[[], [], []], [[], [], []], [[], [], []]]
    b = [[0] * 3, [0] * 3]
    #Ikura Bank
    for tree in xml.findall("./C1/C0/[@Name='Objs']/C1/A0/[@StringValue='Obj_CoopIkuraBank']/.."):
        pos = tree.findall("./C1/[@Name='Translate']/")
        layer = tree.find("./A0/[@Name='LayerConfigName']")
        mX = float(pos[0].attrib["StringValue"])
        mY = float(pos[2].attrib["StringValue"])
        mZ = float(pos[1].attrib["StringValue"])
        if layer.attrib["StringValue"] == "CoopWater_0":
            b[0][0] = mX
            b[0][1] = mY
            b[0][2] = mZ
        if layer.attrib["StringValue"] == "Cmn":
            b[1][0] = mX
            b[1][1] = mY
            b[1][2] = mZ
        plt.scatter(mX, mY, c="black", marker="s")
    for tree in xml.findall("./C1/C0/[@Name='Objs']/C1/A0/[@StringValue='Obj_CoopIkuraBankBase']/.."):
        pos = tree.findall("./C1/[@Name='Translate']/")
        layer = tree.find("./A0/[@Name='LayerConfigName']")
        mX = float(pos[0].attrib["StringValue"])
        mY = float(pos[2].attrib["StringValue"])
        mZ = float(pos[1].attrib["StringValue"])
        plt.scatter(mX, mY, c="black", marker="s")
    for tree in xml.findall("./C1/C0/C1/A0/[@StringValue='%s']/.." % xpath):
        pos = tree.findall("./C1/[@Name='Translate']/")
        layer = tree.find("./A0/[@Name='LayerConfigName']")
        mX = float(pos[0].attrib["StringValue"])
        mY = float(pos[2].attrib["StringValue"])
        mZ = float(pos[1].attrib["StringValue"])

        # Drizzler
        if xpath == "Obj_CoopJumpPointEnemyRocket":
            # if tree.find("./D0/[@Name='IsLinkDest']").attrib["StringValue"] == "False":
            #     continue
            if layer.attrib["StringValue"][:9] != "CoopWater":
                print(layer.attrib["StringValue"][:9])
                # continue
            i = int(layer.attrib["StringValue"][-1])
            if i == 2:
                print(mX, mY, mZ)
            x[i][0].append(mX)
            y[i][0].append(mY)
            z[i][0].append(mZ)
            continue
        # Tower Flyfish Salmonids
        if (xpath[-2:] == "_1") or (xpath[-2:] == "_2") or (xpath[-2:] == "_3"):
            i = int(layer.attrib["StringValue"][-1])
            j = int(xpath[-1]) - 1
            x[i][j].append(mX)
            y[i][j].append(mY)
            z[i][j].append(mZ)
            continue
        # Otherwise
        if pos == None:
            continue
        # print(xpath, layer.attrib["StringValue"])

        if layer.attrib["StringValue"][:9] == "CoopWater" or layer.attrib["StringValue"] == "Cmn":
            if layer.attrib["StringValue"][:9] == "CoopWater":
                i = int(layer.attrib["StringValue"][-1])
            else:
                i = 0
            x[i][0].append(mX)
            y[i][0].append(mY)
            z[i][0].append(mZ)
    if (xpath[-2:] == "_1") or (xpath[-2:] == "_2") or (xpath[-2:] == "_3"):
        i = int(xpath[-1]) - 1
        j = int(layer.attrib["StringValue"][-1])
    # print(xpath, "Size", len(x[0][0]), len(x[1][0]), len(x[2][0]))
    plt.xlim([700, -960])
    # plt.xlim([-960, 700])
    plt.ylim([-600, 500])

    # Display Distance from IkuraBank
    if "Point" in xpath:
        for i in range (0, len(x[0])):
            # print(xpath, b[0], x[0][i], y[0][i], z[0][i])
            for j in range (0, len(x[0][i])): # Low Tide
                dist = math.sqrt(math.pow((b[0][0] - x[0][i][j]), 2) +  math.pow((b[0][1] - y[0][i][j]), 2) + math.pow((b[0][2] - z[0][i][j]), 2))
                plt.annotate(round(dist), (x[0][i][j], y[0][i][j]), color="black")
        for i in range (0, len(x[1])): # Middle Tide
            for j in range (0, len(x[1][i])):
                dist = math.sqrt(math.pow((b[1][0] - x[1][i][j]), 2) +  math.pow((b[1][1] - y[1][i][j]), 2) + math.pow((b[1][2] - z[1][i][j]), 2))
                if xpath != "Obj_CoopJumpPointEnemyRocket":
                    plt.annotate(round(dist), (x[1][i][j], y[1][i][j]), color="black")
        for i in range (0, len(x[2])): # High Tide
            for j in range (0, len(x[2][i])):
                dist = math.sqrt(math.pow((b[1][0] - x[2][i][j]), 2) +  math.pow((b[1][1] - y[2][i][j]), 2) + math.pow((b[1][2] - z[2][i][j]), 2))
                if xpath != "Obj_CoopJumpPointEnemyRocket":
                    plt.annotate(round(dist), (x[2][i][j], y[2][i][j]), color="black")

    # plt.axes().set_aspect('equal', 'datalim')
    plt.scatter(x[0][0], y[0][0], c="red", marker="o")
    plt.scatter(x[0][1], y[0][1], c="red", marker="^")
    plt.scatter(x[0][2], y[0][2], c="red", marker="D")
    plt.scatter(x[1][0], y[1][0], c="blue", marker="o")
    plt.scatter(x[1][1], y[1][1], c="blue", marker="^")
    plt.scatter(x[1][2], y[1][2], c="blue", marker="D")
    plt.scatter(x[2][0], y[2][0], c="yellow", marker="o")
    plt.scatter(x[2][1], y[2][1], c="yellow", marker="^")
    plt.scatter(x[2][2], y[2][2], c="yellow", marker="D")
    if (xpath[-2:] != "_1") and (xpath[-2:] != "_2"):
        # plt.legend()
        plt.tick_params(labelbottom=False,
                    labelleft=False,
                    labelright=False,
                    labeltop=False)
        plt.savefig(str(id) + "_" + xpath +".png", dpi=150, transparent=False, bbox_inches="tight", pad_inches=0.0)
        plt.close()

if __name__ == "__main__":
    print("0: Spawning Grounds(Shakeup)\n1: Marooner's Bay(Shakeship)\n2: Lost Outpost(Shakehouse)\n3: Salmonid Smokeyard(Shakelift)\n4: Ruins of Ark Polaris(Shakeride)\n5: All\nInput the key :")
    id = input()
    if id == "5":
        for id in range(0, 5):
            xmlp = et.XMLParser(encoding="utf-8")
            xml = et.parse("./xml/" + str(id) + ".xml", parser=xmlp)
            config = []
            sum = 0
            for obj in xml.findall("./C1/C0/C1/A0/[@Name='UnitConfigName']"):
                config.append(obj.attrib["StringValue"])
            config = sorted(list(set(config)))
            for xpath in config:
                print("Object Name:", xpath)
                param = xml.findall("./C1/C0/C1/A0/[@StringValue='%s']/.." % xpath)
                sum += len(param)
                print(xpath, len(param))
                loadPos(xml, xpath, id)
            print("UnitConfigName", sum)
        exit(0)
    print("0: Zako SpawnPoint\n1: Boss SpawnPoint\n2: Tower ArrivalPoint\n3: Drizzler JumpPoint\n4: Flyfish ArrivalPoint\n5: All\nInput tht key :")
    type = input()
    print("0: Transparent\n1: None\nInput tht key :")
    alpha = input()
    xmlp = et.XMLParser(encoding="utf-8")
    xml = et.parse("./xml/" + str(id) + ".xml", parser=xmlp)
    loadPos(xml, type, id, alpha)


# Créditos: código del equipo "Rodi Balboa", ganador del torneo de robot-sumo (RoDI)

import rodi
import time
import keyboard

robot = rodi.RoDI()
# for i in range(4):
#   robot.move_forward()
#  time.sleep(1.5)
# robot.move_stop()
# robot.move_left()
# time.sleep(0.55)
# robot.move_stop()
# RODI CANTA
doremi_hz = ["2093", "2349", "2637", "2794", "3136", "3520", "3951", "4699", "4186"]
do = doremi_hz[0]
re = doremi_hz[1]
mi = doremi_hz[2]
fa = doremi_hz[3]
sol = doremi_hz[4]
la = doremi_hz[5]
si = doremi_hz[6]
re1 = doremi_hz[7]
do1 = doremi_hz[8]
cumple = [
    sol,
    sol,
    la,
    sol,
    do,
    si,
    sol,
    sol,
    la,
    sol,
    re1,
    do1,
    do,
    mi,
    sol,
    mi,
    do,
    si,
    la,
    fa,
    fa,
    mi,
    do,
    re,
    do,
]
rodi_entrada = [mi, sol, sol, la, la, la, la]
# print(len(cumple))
# RODI ESQUIVA AUTOMATICO
# for i in range(4):
#   print(robot.see())
#  time.sleep(0.1)
# def ataque():
#    if int(robot.see()) < 10:
#        robot.move_stop()
#        robot.move_left()
#        time.sleep(0.5)

# while True:
#    robot.move_forward()
#    print(robot.see())
#    ataque()
#    time.sleep(0.2)
"""def al_frente():
  if keyboard.is_pressed('w'):
    robot.move_forward()
def izquierda():
  if keyboard.is_pressed('a'):
    robot.move_left()"""

def entrada():
    robot.sing(mi, 250)
    time.sleep(0.25)
    robot.sing(sol, 500)
    time.sleep(0.5)
    robot.sing(la, 2000)
    time.sleep(2)
    robot.sing(la, 250)
    time.sleep(0.25)
    robot.sing(si, 500)
    time.sleep(0.5)
    robot.sing(mi, 2000)
    time.sleep(2)
    # robot.move_forward()
    """for i in range(len(rodi_entrada)):
    robot.sing(int(rodi_entrada[i]),250)
    time.sleep(0.250)
  robot.move_stop()"""

# entrada()

while True:
    if keyboard.is_pressed("w"):
        robot.move_forward()
    elif keyboard.is_pressed("a"):
        robot.move_left()
    elif keyboard.is_pressed("d"):
        robot.move_right()
    elif keyboard.is_pressed("s"):
        robot.move_backward()
    elif keyboard.is_pressed("p"):
        entrada()
    else:
        robot.move_stop()


#!/usr/bin/env python
# coding: utf-8

import pandas as pd
import os
from datetime import datetime

def process_fcast_data(iso_country_code, adm, flnm):
    """
    Process dengue data along with Google Trends and Meteorological data.

    Parameters:
    - iso_country_code (str): ISO country code
    - adm (str): Administrative region
    - flnm (str): Filename of the dengue data. The file should be saved in the 'Dengue' subfolder in 'data'.
                 It should have the following columns: [loc = subnational area name (all caps),
                 cases = number of dengue cases, deaths = number of dengue deaths, Region = regional area name (all caps)]

    Returns:
    - None: Saves the processed data to an Excel file
    """

    # Administrative Boundaries Mapping
    admin_file = os.path.join(os.getcwd(), '..', 'data', f"Admin_Boundaries_{iso_country_code}.csv")
    admin_df = pd.read_csv(admin_file)
    region = admin_df[admin_df["Subregion"] == adm]["Region"].tolist()[0]

    # Google Trends Data
    gt_file = os.path.join(os.getcwd(), '..', 'data', 'Google Trends', f"{iso_country_code}.csv")
    gt_df = pd.read_csv(gt_file)
    gt_df = gt_df[gt_df["region"] == region]

    # Determine the date format
    sample_date = gt_df["date"].iloc[0]
    if "-" in sample_date:
        date_format = '%Y-%m-%d'
    elif "/" in sample_date:
        date_format = '%d/%m/%Y'
    else:
        print("Unknown date format in Google Trends data.")
        return

    gt_df["date"] = gt_df["date"].apply(lambda d: datetime.strptime(d, '%Y-%m-%d'))
    gt_df = gt_df.pivot(index=["date"], columns="keyword", values="value").reset_index().sort_values(by=["date"], ascending=True)

    # Meteorological Data
    weather_file = os.path.join(os.getcwd(), '..', 'data', 'NASA POWER', f"{iso_country_code}.csv")
    weather_df = pd.read_csv(weather_file)
    weather_df = weather_df[weather_df["loc"] == adm.upper()]

    # Check if DataFrame is empty
    if weather_df.empty:
        print(f"No weather data found for {adm.upper()}. Skipping...")
        return

    # Check if required columns exist
    if 'YEAR' not in weather_df.columns or 'DOY' not in weather_df.columns:
        print("Required columns 'YEAR' and 'DOY' not found in weather data. Skipping...")
        return

    # If everything is fine, proceed with the date conversion
    weather_df['date'] = weather_df.apply(lambda row: datetime.strptime(f"{int(row['YEAR'])} {int(row['DOY'])}", '%Y %j').strftime('%d/%m/%Y'), axis=1)
    weather_df["date"] = weather_df["date"].apply(lambda d: datetime.strptime(d, '%d/%m/%Y'))
    weather_df = weather_df.drop(columns=["YEAR","DOY"])

    # Combine Google Trends and Meteorological Data for the specific location
    combin_df = gt_df.merge(weather_df, on="date", how="left")

    # Dengue Data
    dengue_path = os.path.join(os.getcwd(), '..', 'data/Dengue', f"{iso_country_code}/", flnm)
    file_name, file_extension = os.path.splitext(dengue_path)

    if file_extension == ".xlsx":
        dengue_df = pd.read_excel(dengue_path)
    elif file_extension == ".csv":
        dengue_df = pd.read_csv(dengue_path)
    else:
        print("Wrong file type. Read either a CSV or an Excel file")
        return

    dengue_df = dengue_df[dengue_df["loc"] == adm.upper()]
    dengue_df = dengue_df.merge(combin_df, on=["loc","date"], how="left")

    # Save the processed data
    data_folder = os.path.join(os.getcwd(), '..', 'processed', 'Forecasting', f"{iso_country_code}")
    os.makedirs(data_folder, exist_ok=True)
    flnm2 = os.path.join(data_folder, f"{adm}.xlsx")
    dengue_df.to_excel(flnm2, index=False)
    print(f"Finished with {adm}")

# Example usage
#iso_country_code = "PHL"
#admin_file = os.path.join(os.getcwd(), '..', 'data', f"Admin_Boundaries_{iso_country_code}.csv")
#admin_df = pd.read_csv(admin_file)
#adm_list = admin_df["Subregion"].tolist()
#for a in adm_list:
#    process_fcast_data(iso_country_code, a, f"{iso_country_code}.xlsx")


from azure.cognitiveservices.vision.computervision import ComputerVisionClient
from utils import authenticate
import pandas as pd

class BaseVision:
    def __init__(
        self, service_endpoint: str, service_key: str = None, cognitive_key: str = None
    ) -> None:

        self.vision_key = service_key
        self.url = service_endpoint
        credentials = authenticate(service_key, cognitive_key)
        self.service_client = ComputerVisionClient(self.url, credentials)

    def get_metadata(self, response):
        """
        Collects metadata for each image
        in the response object to pandas.data.frame
        """

        tmp = {k: pd.DataFrame(v["metadata"], index=[0]) for k, v in response.items()}
        for k, v in tmp.items():
            tmp[k]["image"] = k
            tmp[k]["request_id"] = v["request_id"]

        metadata = pd.concat(tmp.values(), ignore_index=True, sort=False)

        return metadata


import numpy as np
import random
import math
from collections import namedtuple
from itertools import count

import torch
import torch.optim as optim
import torch.nn as nn

from buffers.dqn_replay_memory import ReplayMemory
from models.dqn_model import DQN
from models.dqn_resnet_model import ResNet
# from models.dqn_multimodal_resnet_model import MultimodalCNN
from models.dqn_multimodal_resnet_model import MultimodalResnetAndFC
from agents.agent import Agent
from utils.plotting import plot_rewards

class DQNAgent(Agent):
    """
    Source for most code before it was rearranged/changed, and details:
    https://pytorch.org/tutorials/intermediate/reinforcement_q_learning.html
    """
    def __init__(self, env, batch_size=128, gamma=0.99, eps_start=0.9, eps_end=0.05,
                 eps_decay=1000, tau=0.005, lr=1e-4, replay_mem_size=10000, model_type="fc", use_action_mask=False):
        super().__init__(env)
        # constants
        self.batch_size = batch_size  # the number of transitions sampled from the replay buffer
        self.gamma = gamma  # the discount factor
        self.eps_start = eps_start  # the starting value of epsilon
        self.eps_end = eps_end  # the final value of epsilon
        self.eps_decay = eps_decay  # controls the rate of exponential decay of epsilon, higher means a slower decay
        self.tau = tau  # the update rate of the target network
        self.use_action_mask = use_action_mask  # to use action mask or not (helps computational efficiency when not)

        # torch, networks
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        observation, _ = self.env.reset()
        n_actions = self.env.action_space.n

        if model_type == "fc":
            # get values for policy/target net dims
            n_observations = len(observation)
            self.policy_net = DQN(n_observations, n_actions).to(self.device)
            self.target_net = DQN(n_observations, n_actions).to(self.device)
        elif model_type == "resnet":
            self.policy_net = ResNet(2, observation.shape, n_actions).to(self.device)
            self.target_net = ResNet(2, observation.shape, n_actions).to(self.device)
        elif model_type == "multires":
            input_shapes = {modality: array.shape for modality, array in observation.items()}
            self.policy_net = MultimodalResnetAndFC(0, input_shapes, 32, n_actions).to(self.device)
            self.target_net = MultimodalResnetAndFC(0, input_shapes, 32, n_actions).to(self.device)
            # self.policy_net = MultimodalCNN(2, input_shapes, n_actions).to(self.device)
            # self.target_net = MultimodalCNN(2, input_shapes, n_actions).to(self.device)
        self.target_net.load_state_dict(self.policy_net.state_dict())
        self.target_net.eval()
        self.optimizer = optim.AdamW(self.policy_net.parameters(), lr=lr, amsgrad=True)

        # replay memory
        self.Transition = namedtuple('Transition',
                                     ('state', 'action', 'next_state', 'next_action_mask', 'reward'))
        self.memory = ReplayMemory(replay_mem_size, self.Transition, self.device)

    def select_action(self, obs, action_mask=None) -> int:
        self.policy_net.eval()
        return int(self.select_action_with_eps(obs, self.eps_end, action_mask))

    def select_action_with_eps(self, obs: np.ndarray, eps_threshold, action_mask=None) -> int:
        self.policy_net.eval()
        sample = random.random()
        if sample > eps_threshold:
            obs = torch.tensor(obs, dtype=torch.float32, device=self.device).unsqueeze(0)
            with torch.no_grad():
                # t.max(1) will return the largest column value of each row.
                # second column on max result is index of where max element was
                # found, so we pick action with the larger expected reward.
                if action_mask is not None:
                    # if we have an action mask, only consider those actions
                    # so get which index in action_mask? the one with the highest value, as it was taken from policy
                    return int(action_mask[self.policy_net(obs)[:, action_mask].max(1)[1].view(1, 1).item()])
                else:
                    return self.policy_net(obs).max(1)[1].view(1, 1).item()
        else:
            # cast to int because sample() returns np.int64, for consistency
            if action_mask is not None:
                return int(np.random.choice(action_mask, 1)[0])
            else:
                return int(self.env.action_space.sample())

    def train(self, epochs: int, early_stopping_rounds: int = -1, early_stopping_threshold: float = 200.0,
              show_progress: bool = False, print_progress: int = 0, save_checkpoints: int = 0) -> list:
        steps_done = 0
        epoch_rewards = []
        for i_episode in range(epochs):
            # Initialize the environment and get it's state
            total_reward = 0
            observation, info = self.env.reset()
            action_mask = info['action_mask'] if 'action_mask' in info else None
            # observation_tensor = torch.tensor(observation, dtype=torch.float32, device=self.device).unsqueeze(0)
            for t in count():
                current_eps = self.eps_end + (self.eps_start - self.eps_end) * \
                                math.exp(-1. * steps_done / self.eps_decay)
                action = self.select_action_with_eps(observation, current_eps, action_mask)
                action_tensor = torch.tensor([[action]], device=self.device, dtype=torch.long)
                next_observation, reward, terminated, truncated, next_info = self.env.step(action)
                next_action_mask = next_info['action_mask'] if 'action_mask' in next_info else None
                steps_done += 1
                total_reward += reward
                reward_tensor = torch.tensor([reward], device=self.device)
                done = terminated or truncated

                if terminated:
                    next_observation = None
                    next_action_mask = None
                # else:
                #     next_observation_tensor = torch.tensor(next_observation,
                #                                            dtype=torch.float32,
                #                                            device=self.device).unsqueeze(0)

                # Store the transition in memory
                self.memory.push(observation, action_tensor, next_observation, next_action_mask, reward_tensor)

                # Move to the next state
                observation = next_observation
                action_mask = next_action_mask

                # observation_tensor = next_observation_tensor

                # Perform one step of the optimization (on the policy network)
                self.optimize_model()

                # Soft update of the target network's weights
                # θ′ ← τ θ + (1 −τ )θ′
                target_net_state_dict = self.target_net.state_dict()
                policy_net_state_dict = self.policy_net.state_dict()
                for key in policy_net_state_dict:
                    target_net_state_dict[key] = policy_net_state_dict[key]*self.tau + target_net_state_dict[key]*(1-self.tau)
                self.target_net.load_state_dict(target_net_state_dict)

                if done:
                    epoch_rewards.append(total_reward)
                    # if it is positive, we want to stop early (when it reaches some threshold)
                    if early_stopping_rounds and len(epoch_rewards) >= early_stopping_rounds:
                        if (sum(epoch_rewards[-early_stopping_rounds:]) / early_stopping_rounds) > early_stopping_threshold:
                            return epoch_rewards

                    if show_progress:
                        plot_rewards(epoch_rewards)
                    if print_progress > 0 and i_episode % print_progress == 0:
                        print(f"episode: {i_episode}, reward: {total_reward}, eps: {current_eps}")
                    if save_checkpoints > 0 and i_episode % save_checkpoints == 0:
                        self.save_model("current_dqn_training")
                    break

        return epoch_rewards

    def optimize_model(self) -> None:
        # we train the policy net here
        self.policy_net.train()

        if len(self.memory) < self.batch_size:
            return
        transitions = self.memory.sample(self.batch_size)
        # Transpose the batch (see https://stackoverflow.com/a/19343/3343043 for
        # detailed explanation). This converts batch-array of Transitions
        # to Transition of batch-arrays.
        batch = self.Transition(*zip(*transitions))

        # Compute a mask of non-final states and concatenate the batch elements
        # (a final state would've been the one after which simulation ended)
        non_final_mask = torch.tensor(tuple(map(lambda s: s is not None,
                                                batch.next_state)), device=self.device, dtype=torch.bool)
        non_final_next_states = torch.cat([s for s in batch.next_state
                                           if s is not None])
        state_batch = torch.cat(batch.state)
        action_batch = torch.cat(batch.action)
        reward_batch = torch.cat(batch.reward)
        if self.use_action_mask:
            # Pad the mask arrays to the same length
            max_len = max(len(row_mask) for row_mask in batch.next_action_mask if row_mask is not None)
            # note: instead of padding with zero, could also just pad with row_mask[0] which is some (first) action
            # in the mask, because when we take max anyway, doesn't matter if it is duplicated. But 0 works for
            # solitaire
            next_action_mask_batch = [np.append(row_mask, [row_mask[0]] * (max_len - len(row_mask))) for row_mask in batch.next_action_mask if row_mask is not None]
            # next_action_mask_batch = [np.append(row_mask, [0] * (max_len - len(row_mask))) for row_mask in batch.next_action_mask if row_mask is not None]
            next_action_mask_batch = torch.tensor(np.array(next_action_mask_batch), dtype=torch.long)

        # Compute Q(s_t, a) - the model computes Q(s_t), then we select the
        # columns of actions taken. These are the actions which would've been taken
        # for each batch state according to policy_net
        state_action_values = self.policy_net(state_batch).gather(1, action_batch)

        # Compute V(s_{t+1}) for all next states.
        # Expected values of actions for non_final_next_states are computed based
        # on the "older" target_net; selecting their best reward with max(1)[0].
        # This is merged based on the mask, such that we'll have either the expected
        # state value or 0 in case the state was final.
        next_state_values = torch.zeros(self.batch_size, device=self.device)
        with torch.no_grad():
            if self.use_action_mask:
                # https://discuss.pytorch.org/t/selecting-from-a-2d-tensor-with-rows-of-column-indexes/167717
                next_state_values[non_final_mask] = self.target_net(non_final_next_states)[torch.arange(len(non_final_next_states)).unsqueeze(1), next_action_mask_batch].max(1)[0]
            else:
                next_state_values[non_final_mask] = self.target_net(non_final_next_states).max(1)[0]
        # Compute the expected Q values
        expected_state_action_values = (next_state_values * self.gamma) + reward_batch

        # Compute Huber loss
        criterion = nn.SmoothL1Loss()
        loss = criterion(state_action_values, expected_state_action_values.unsqueeze(1))

        # Optimize the model
        self.optimizer.zero_grad()
        loss.backward()
        # In-place gradient clipping, don't want to do norm because we have sparse gradients and
        # we don't want norm to penalize as a whole just the individuals
        torch.nn.utils.clip_grad_value_(self.policy_net.parameters(), 10)
        self.optimizer.step()

    def eval(self, num_episodes: int) -> list:
        return super().eval(num_episodes)

    def investigate_model_outputs(self, obs: np.ndarray) -> np.ndarray:
        obs = torch.tensor(obs, dtype=torch.float32, device=self.device).unsqueeze(0)
        with torch.no_grad():
            # t.max(1) will return the largest column value of each row.
            # second column on max result is index of where max element was
            # found, so we pick action with the larger expected reward.
            return self.policy_net(obs).detach().cpu().numpy()[0]

    # Save the agent's model or parameters to a file
    def save_model(self, filepath) -> None:
        print(f"saving model to {filepath}.pt")
        filepath_with_extension = filepath + ".pt"
        torch.save(self.policy_net.state_dict(), filepath_with_extension)

    # Load the agent's model or parameters from a file
    def load_model(self, filepath) -> None:
        print(f"loading model {filepath}")
        self.policy_net.load_state_dict(torch.load(filepath, map_location=torch.device(self.device)))


import json
import structlog
import pika
import config
from multiprocessing import Process
from ServiceManager import ServiceInfo, cleanupService

logger = structlog.get_logger()

def cleanupCallback(ch, method, properties, body):
    """Take service off of queue and delete its corresponding container"""

    logger.info("cleanupCallback", msg="Recieved Message", body=body)
    body = json.loads(body)
    service = ServiceInfo(body)
    log = logger.bind(service=service.__dict__)
    log.info("cleanupCallback", msg="Cleaning Up Service")
    cleanupService(service)
    ch.basic_ack(delivery_tag=method.delivery_tag)
    log.info("cleanupCallback", msg="Service Cleanup Done")

def cleanupWorker():
    """Declare cleanup queue and callback"""

    connection = pika.BlockingConnection(pika.ConnectionParameters(config.RABBITMQ_SERVER))
    channel = connection.channel()
    channel.queue_declare(queue='cleanupQueue', durable=True)
    logger.info("cleanupWorker", msg="Starting Cleanup Worker", queue="cleanupQueue")
    channel.basic_consume(cleanupCallback, queue='cleanupQueue')
    channel.start_consuming()

def startCleanupWorkers(numThreads):
    """Start numThreads cleanup workers"""
    for i in range(numThreads):
        t = Process(target=cleanupWorker)
        t.daemon = True
        t.start()


from enum import Enum, IntEnum

class Level(Enum):
    """
    Error levels.
    """

    ERROR = "E"
    WARNING = "W"
    INFO = "I"

    def __str__(self):
        return self.value

class PoolMaintenanceErrorCode(IntEnum):
    """
    Maintenance error codes for the pool.
    """

    NO_IPC_REQUESTS = 1
    NO_POOL_CHANGES = 2

    def __str__(self):
        return f"{Level.ERROR}M{str(self.value).zfill(3)}"

    @staticmethod
    def from_str(code_str):
        """
        Discover the code, if any, from the code string.

        :returns: the code if it finds a match, otherwise None
        :rtype: PoolMaintenanceErrorCode or NoneType
        """
        return next(
            (code for code in PoolMaintenanceErrorCode if code_str == str(code)), None
        )

    def explain(self):
        """
        Return an explanation of the error return code.
        """
        if self is PoolMaintenanceErrorCode.NO_IPC_REQUESTS:
            return (
                "The pool will return an error on any IPC request that could "
                "cause a change in the pool state, for example, a request to "
                "rename a filesystem. It will still be able to respond to "
                "purely informational requests."
            )

        if self is PoolMaintenanceErrorCode.NO_POOL_CHANGES:
            return (
                "The pool is unable to manage itself by reacting to events, "
                "such as devicemapper events, that might require it to take "
                "any maintenance operations."
            )

        assert False, "impossible error code reached"  # pragma: no cover

    def summarize(self):
        """
        Return a short summary of the return code.
        """
        if self is PoolMaintenanceErrorCode.NO_IPC_REQUESTS:
            return "Pool state changes not possible"

        if self is PoolMaintenanceErrorCode.NO_POOL_CHANGES:
            return "Pool maintenance operations not possible"

        assert False, "impossible error code reached"  # pragma: no cover

class PoolAllocSpaceErrorCode(IntEnum):
    """
    Code if the pool has run out of space to allocate.
    """

    NO_ALLOC_SPACE = 1

    def __str__(self):
        return f"{Level.WARNING}S{str(self.value).zfill(3)}"

    def explain(self):
        """
        Return an explanation of the return code.
        """
        if self is PoolAllocSpaceErrorCode.NO_ALLOC_SPACE:
            return (
                "Every device belonging to the pool has been fully allocated. "
                "To increase the allocable space, add additional data devices "
                "to the pool."
            )

        assert False, "impossible error code reached"  # pragma: no cover

    def summarize(self):
        """
        Return a short summary of the return code.
        """
        if self is PoolAllocSpaceErrorCode.NO_ALLOC_SPACE:
            return "All devices fully allocated"

        assert False, "impossible error code reached"  # pragma: no cover

    @staticmethod
    def from_str(code_str):
        """
        Discover the code, if any, from the code string.

        :returns: the code if it finds a match, otherwise None
        :rtype: PoolAllocSpaceErrorCode or NoneType
        """
        return next(
            (code for code in PoolAllocSpaceErrorCode if code_str == str(code)), None
        )

class PoolDeviceSizeChangeCode(IntEnum):
    """
    Codes for identifying for a pool if a device that belongs to the pool has
    been detected to have increased or reduced in size.
    """

    DEVICE_SIZE_INCREASED = 1
    DEVICE_SIZE_DECREASED = 2

    def __str__(self):
        if self is PoolDeviceSizeChangeCode.DEVICE_SIZE_INCREASED:
            return f"{Level.INFO}DS{str(self.value).zfill(3)}"

        if self is PoolDeviceSizeChangeCode.DEVICE_SIZE_DECREASED:
            return f"{Level.WARNING}DS{str(self.value).zfill(3)}"

        assert False, "impossible error code reached"  # pragma: no cover

    def explain(self):
        """
        Return an explanation of the return code.
        """
        if self is PoolDeviceSizeChangeCode.DEVICE_SIZE_INCREASED:
            return (
                "At least one device belonging to this pool appears to have "
                "increased in size."
            )

        if self is PoolDeviceSizeChangeCode.DEVICE_SIZE_DECREASED:
            return (
                "At least one device belonging to this pool appears to have "
                "decreased in size."
            )

        assert False, "impossible error code reached"  # pragma: no cover

    def summarize(self):
        """
        Return a short summary of the return code.
        """
        if self is PoolDeviceSizeChangeCode.DEVICE_SIZE_INCREASED:
            return "A device in this pool has increased in size."

        if self is PoolDeviceSizeChangeCode.DEVICE_SIZE_DECREASED:
            return "A device in this pool has decreased in size."

        assert False, "impossible error code reached"  # pragma: no cover

    @staticmethod
    def from_str(code_str):
        """
        Discover the code, if any, from the code string.

        :returns: the code if it finds a match, otherwise None
        :rtype: PoolAllocSpaceErrorCode or NoneType
        """
        return next(
            (code for code in PoolDeviceSizeChangeCode if code_str == str(code)), None
        )

class PoolErrorCode:
    """
    Summary class for all pool error codes.
    """

    CLASSES = [
        PoolMaintenanceErrorCode,
        PoolAllocSpaceErrorCode,
        PoolDeviceSizeChangeCode,
    ]

    @staticmethod
    def codes():
        """
        Return all pool error codes.
        """
        return [code for c in PoolErrorCode.CLASSES for code in list(c)]

    @staticmethod
    def error_from_str(error_code):
        """
        Obtain an error object from a distinguishing error string.

        :param str error_code:
        :returns: error object
        :raises: StopIteration if no match found
        """
        return next(
            (
                code
                for code in (c.from_str(error_code) for c in PoolErrorCode.CLASSES)
                if code is not None
            )
        )

    @staticmethod
    def explain(error_code):
        """
        Return explanation for error code, else None.

        :param str error_code:
        """
        return PoolErrorCode.error_from_str(error_code).explain()