import logging
from collections import defaultdict
from typing import Dict, Optional, Tuple, Union

import numpy as np
from matplotlib import rcParams
from matplotlib.axes import SubplotBase
from matplotlib.axis import Axis
from matplotlib.colors import LogNorm
from matplotlib.ticker import AutoMinorLocator, MaxNLocator
from pymatgen.electronic_structure.bandstructure import (
    BandStructure,
    BandStructureSymmLine,
)
from pymatgen.electronic_structure.core import Spin
from pymatgen.electronic_structure.plotter import BSPlotter
from sumo.plotting import pretty_plot
from sumo.symmetry import Kpath, PymatgenKpath

from amset.constants import defaults, hbar
from amset.interpolation.bandstructure import Interpolator
from amset.interpolation.periodic import PeriodicLinearInterpolator
from amset.log import initialize_amset_logger
from amset.plot import BaseMeshPlotter, amset_base_style, styled_plot

__author__ = "Alex Ganose"
__maintainer__ = "Alex Ganose"
__email__ = "aganose@lbl.gov"

logger = logging.getLogger(__name__)

class LineshapePlotter(BaseMeshPlotter):
    def __init__(
        self,
        data,
        interpolation_factor=5,
        print_log=defaults["print_log"],
        symprec=defaults["symprec"],
    ):
        super().__init__(data)
        self.interpolation_factor = interpolation_factor

        if print_log:
            initialize_amset_logger(filename="lineshape.log")
        self.symprec = symprec

    def _get_interpolater(self, n_idx, t_idx, mode="linear"):
        props = defaultdict(dict)
        for spin in self.spins:
            # calculate total rate
            spin_rates = np.sum(self.scattering_rates[spin][:, n_idx, t_idx], axis=0)

            # easier to interpolate the log
            log_rates = np.log10(spin_rates)

            # # handle rates that close to numerical noise
            log_rates[log_rates > 18] = 15
            log_rates[np.isnan(log_rates)] = 15

            # map to full k-point mesh
            props[spin]["rates"] = log_rates

        if mode == "linear":
            return _LinearBandStructureInterpolator(
                self.kpoints,
                self.ir_to_full_kpoint_mapping,
                self.energies,
                self.structure,
                self.efermi,
                props,
            )
        elif mode == "fourier":
            bs = BandStructure(
                self.ir_kpoints,
                self.energies,
                self.structure.lattice,
                self.efermi,
                structure=self.structure,
            )

            return Interpolator(
                bs,
                self.num_electrons,
                interpolation_factor=self.interpolation_factor,
                soc=self.soc,
                other_properties=props,
            )
        raise ValueError("Unknown interpolation mode; should be 'linear' or 'fourier'.")

    @styled_plot(amset_base_style)
    def get_plot(
        self,
        n_idx,
        t_idx,
        zero_to_efermi=True,
        estep=0.01,
        line_density=100,
        height=3.2,
        width=3.2,
        emin=None,
        emax=None,
        amin=5e-5,
        amax=1e-1,
        ylabel="Energy (eV)",
        plt=None,
        aspect=None,
        kpath=None,
        cmap="viridis",
        colorbar=True,
        style=None,
        no_base_style=False,
        fonts=None,
    ):
        interpolater = self._get_interpolater(n_idx, t_idx)

        bs, prop = interpolater.get_line_mode_band_structure(
            line_density=line_density,
            return_other_properties=True,
            kpath=kpath,
            symprec=self.symprec,
        )
        bs, rates = force_branches(bs, {s: p["rates"] for s, p in prop.items()})

        fd_emin, fd_emax = self.fd_cutoffs
        if not emin:
            emin = fd_emin
            if zero_to_efermi:
                emin -= bs.efermi

        if not emax:
            emax = fd_emax
            if zero_to_efermi:
                emax -= bs.efermi

        logger.info("Plotting band structure")
        if isinstance(plt, (Axis, SubplotBase)):
            ax = plt
        else:
            plt = pretty_plot(width=width, height=height, plt=plt)
            ax = plt.gca()

        if zero_to_efermi:
            bs.bands = {s: b - bs.efermi for s, b in bs.bands.items()}
            bs.efermi = 0

        bs_plotter = BSPlotter(bs)
        plot_data = bs_plotter.bs_plot_data(zero_to_efermi=zero_to_efermi)

        energies = np.linspace(emin, emax, int((emax - emin) / estep))
        distances = np.array([d for x in plot_data["distances"] for d in x])

        # rates are currently log(rate)
        mesh_data = np.full((len(distances), len(energies)), 0.0)
        for spin in self.spins:
            for spin_energies, spin_rates in zip(bs.bands[spin], rates[spin]):
                for d_idx in range(len(distances)):
                    energy = spin_energies[d_idx]
                    linewidth = 10 ** spin_rates[d_idx] * hbar / 2
                    broadening = lorentzian(energies, energy, linewidth)
                    broadening /= 1000  # convert 1/eV to 1/meV
                    mesh_data[d_idx] += broadening

        im = ax.pcolormesh(
            distances,
            energies,
            mesh_data.T,
            rasterized=True,
            cmap=cmap,
            norm=LogNorm(vmin=amin, vmax=amax),
            shading="auto",
        )
        if colorbar:
            pos = ax.get_position()
            cax = plt.gcf().add_axes([pos.x1 + 0.035, pos.y0, 0.035, pos.height])
            cbar = plt.colorbar(im, cax=cax)
            cbar.ax.tick_params(axis="y", length=rcParams["ytick.major.size"] * 0.5)
            cbar.ax.set_ylabel(
                r"$A_\mathbf{k}$ (meV$^{-1}$)", rotation=270, va="bottom"
            )

        _maketicks(ax, bs_plotter, ylabel=ylabel)
        _makeplot(
            ax,
            plot_data,
            bs,
            zero_to_efermi=zero_to_efermi,
            width=width,
            height=height,
            ymin=emin,
            ymax=emax,
            aspect=aspect,
        )
        return plt

def _makeplot(
    ax,
    data,
    bs,
    zero_to_efermi=True,
    ymin=-3.0,
    ymax=3.0,
    height=None,
    width=None,
    aspect=None,
):
    """Tidy the band structure & add the density of states if required."""
    # draw line at Fermi level if not zeroing to e-Fermi
    if not zero_to_efermi:
        ytick_color = rcParams["ytick.color"]
        ef = bs.efermi
        ax.axhline(ef, color=ytick_color)

    # set x and y limits
    ax.set_xlim(0, data["distances"][-1][-1])
    if bs.is_metal() and not zero_to_efermi:
        ax.set_ylim(bs.efermi + ymin, bs.efermi + ymax)
    else:
        ax.set_ylim(ymin, ymax)

    # keep correct aspect ratio for axes based on canvas size
    if aspect is not False:
        x0, x1 = ax.get_xlim()
        y0, y1 = ax.get_ylim()
        if width is None:
            width = rcParams["figure.figsize"][0]
        if height is None:
            height = rcParams["figure.figsize"][1]

        if not aspect:
            aspect = height / width

        ax.set_aspect(aspect * ((x1 - x0) / (y1 - y0)))

def _maketicks(ax, bs_plotter, ylabel="Energy (eV)"):
    """Utility method to add tick marks to a band structure."""
    # set y-ticks
    ax.yaxis.set_major_locator(MaxNLocator(6))
    ax.yaxis.set_minor_locator(AutoMinorLocator(2))

    # set x-ticks; only plot the unique tick labels
    ticks = bs_plotter.get_ticks()
    unique_d = []
    unique_l = []
    if ticks["distance"]:
        temp_ticks = list(zip(ticks["distance"], ticks["label"]))
        unique_d.append(temp_ticks[0][0])
        unique_l.append(temp_ticks[0][1])
        for i in range(1, len(temp_ticks)):
            # Append label to sequence if it is not same as predecessor
            if unique_l[-1] != temp_ticks[i][1]:
                unique_d.append(temp_ticks[i][0])
                unique_l.append(temp_ticks[i][1])

    logging.info("Label positions:")
    for dist, label in list(zip(unique_d, unique_l)):
        logging.info(f"\t{dist:.4f}: {label}")

    ax.set_xticks(unique_d)
    ax.set_xticklabels(unique_l)
    ax.xaxis.grid(True)
    ax.set_ylabel(ylabel)

def lorentzian(x, x0, gamma):
    return 1 / np.pi * gamma / ((x - x0) ** 2 + gamma**2)

class _LinearBandStructureInterpolator:
    def __init__(
        self,
        full_kpoints,
        ir_to_full_idx,
        energies,
        structure,
        efermi,
        other_properties,
    ):
        self.structure = structure
        self.efermi = efermi
        self.spins = list(energies.keys())
        self.nbands = {s: len(e) for s, e in energies.items()}

        full_energies = {s: e[:, ir_to_full_idx] for s, e in energies.items()}
        self.bs_interpolator = PeriodicLinearInterpolator.from_data(
            full_kpoints, full_energies
        )

        self.property_interpolators = {}
        other_properties = _transpose_dict(other_properties)
        for prop, prop_data in other_properties.items():
            full_prop_data = {s: p[:, ir_to_full_idx] for s, p in prop_data.items()}
            self.property_interpolators[prop] = PeriodicLinearInterpolator.from_data(
                full_kpoints, full_prop_data, gaussian=0.75
            )

    def get_line_mode_band_structure(
        self,
        line_density: int = 50,
        kpath: Optional[Kpath] = None,
        symprec: Optional[float] = defaults["symprec"],
        return_other_properties: bool = False,
    ) -> Union[
        BandStructureSymmLine,
        Tuple[BandStructureSymmLine, Dict[Spin, Dict[str, np.ndarray]]],
    ]:
        """Gets the interpolated band structure along high symmetry directions.

        Args:
            line_density: The maximum number of k-points between each two
                consecutive high-symmetry k-points
            symprec: The symmetry tolerance used to determine the space group
                and high-symmetry path.
            return_other_properties: Whether to include the interpolated
                other_properties data for each k-point along the band structure path.

        Returns:
            The line mode band structure.
        """
        if not kpath:
            kpath = PymatgenKpath(self.structure, symprec=symprec)

        kpoints, labels = kpath.get_kpoints(line_density=line_density, cart_coords=True)
        labels_dict = {
            label: kpoint for kpoint, label in zip(kpoints, labels) if label != ""
        }

        rlat = self.structure.lattice.reciprocal_lattice
        frac_kpoints = rlat.get_fractional_coords(kpoints)

        energies = {}
        other_properties = defaultdict(dict)
        for spin in self.spins:
            energies[spin] = self._interpolate_spin(
                spin, frac_kpoints, self.bs_interpolator
            )

            if return_other_properties:
                for prop, property_interpolator in self.property_interpolators.items():
                    other_properties[spin][prop] = self._interpolate_spin(
                        spin, frac_kpoints, property_interpolator
                    )

        bs = BandStructureSymmLine(
            kpoints,
            energies,
            rlat,
            self.efermi,
            labels_dict,
            coords_are_cartesian=True,
            structure=self.structure,
        )

        if return_other_properties:
            return bs, other_properties
        else:
            return bs

    def _interpolate_spin(self, spin, kpoints, interpolator):
        nkpoints = len(kpoints)
        spin_nbands = self.nbands[spin]
        ibands = np.repeat(np.arange(spin_nbands), nkpoints)
        all_kpoints = np.tile(kpoints, (spin_nbands, 1))
        data = interpolator.interpolate(spin, ibands, all_kpoints)
        return data.reshape(spin_nbands, nkpoints)

def _transpose_dict(d):
    td = defaultdict(dict)
    for k1, v1 in d.items():
        for k2, v2 in v1.items():
            td[k2][k1] = v2
    return td

def force_branches(bandstructure, other_property=None):
    """Force a linemode band structure to contain branches.

    Branches give a specific portion of the path from one high-symmetry point
    to another. Branches are required for the plotting methods to function correctly.
    Unfortunately, due to the pymatgen BandStructure implementation they require
    duplicate k-points in the band structure path. To avoid this unnecessary
    computational expense, this function can reconstruct branches in band structures
    without the duplicate k-points.

    Args:
        bandstructure: A band structure object.
        other_property: Another property with the format {spin: (nbands, nkpts, ...)
            to split into branches.

    Returns:
        A band structure with brnaches.
    """
    kpoints = np.array([k.frac_coords for k in bandstructure.kpoints])
    labels_dict = {k: v.frac_coords for k, v in bandstructure.labels_dict.items()}

    # pymatgen band structure objects support branches. These are formed when
    # two kpoints with the same label are next to each other. This bit of code
    # will ensure that the band structure will contain branches, if it doesn't
    # already.
    dup_ids = []
    high_sym_kpoints = tuple(map(tuple, labels_dict.values()))
    for i, k in enumerate(kpoints):
        dup_ids.append(i)
        if (
            tuple(k) in high_sym_kpoints
            and i != 0
            and i != len(kpoints) - 1
            and (
                not np.array_equal(kpoints[i + 1], k)
                or not np.array_equal(kpoints[i - 1], k)
            )
        ):
            dup_ids.append(i)

    kpoints = kpoints[dup_ids]

    eigenvals = {}
    projections = {}
    for spin, spin_energies in bandstructure.bands.items():
        eigenvals[spin] = spin_energies[:, dup_ids]
        if len(bandstructure.projections) != 0:
            projections[spin] = bandstructure.projections[spin][:, dup_ids]

    new_property = {}
    if other_property is not None:
        for spin, spin_prop in other_property.items():
            new_property[spin] = spin_prop[:, dup_ids]

    new_bandstructure = type(bandstructure)(
        kpoints,
        eigenvals,
        bandstructure.lattice_rec,
        bandstructure.efermi,
        labels_dict,
        structure=bandstructure.structure,
        projections=projections,
    )
    return new_bandstructure, new_property


BOARD_SIZE = 9 # it's a square

def strip_zero_unique(list):
  strip_zero = [i for i in list if i != 0]

  return len(strip_zero) == len(set(strip_zero))

class Board:
  def __init__(self):
    self.cells = [ [ 0 for _ in range(1 + BOARD_SIZE) ] for _ in range(1 + BOARD_SIZE) ]

    self.known_points = []
    self.black_dot_constraints = []
    self.white_dot_constraints = []
    self.arrow_constraints = []
    self.thermo_constraints = []
    self.palindrome_constraints = []
    self.box_constraints = []

  def get_cells(self):
    return [ line[1:] for line in self.cells[1:] ]

  def set_cell(self, line, col, value):
    self.cells[line][col] = value

  def add_known_point(self, line, col, value):
    self.known_points.append((line, col, value))

  # black dot connects two adjacent cells, which must have a ratio of 2:1
  def add_black_dot_constraint(self, c1, c2):
    self.black_dot_constraints.append((c1, c2))

  # white dot connects two adjacent cells, which must contain consecutive numbers
  def add_white_dot_constraint(self, c1, c2):
    self.white_dot_constraints.append((c1, c2))

  # the arrow connects multiple cells, and the numbers must sum up to the start of the arrow
  # arrow will be passed as a list of cells, from start to end
  def add_arrow_constraint(self, cell_list):
    self.arrow_constraints.append(cell_list)

  # thermo connects multiple cells, and the numbers must be increasing from the bottom up
  # thermo will be passed as a list of cells, from bottom to top
  def add_thermo_constraint(self, cell_list):
    self.thermo_constraints.append(cell_list)

  # palindrome is self-explanatory
  # also passed as a list of cells, in the specific order
  def add_palindrome_constraint(self, cell_list):
    self.palindrome_constraints.append(cell_list)

  # box constains multiple cells, and the numbers must be unique and sum to a specific number
  # box will be passed as a tuple containing the sum and the list of cells, in any order
  def add_box_constraint(self, sum, cell_list):
    self.box_constraints.append((sum, cell_list))

  def check_known_points(self, line, col):
    for constraint in self.known_points:
      if constraint[0] == line and constraint[1] == col:
        return constraint[2] == self.cells[line][col]

    return True

  def check_line(self, line, col):
    # print(f'For line {line}, col {col}, checking {self.cells[line][col]} vs {self.cells[line][1:col]}')
    return self.cells[line][col] not in self.cells[line][1:col]

  def check_col(self, line, col):
    # print(f'For line {line}, col {col}, checking {self.cells[line][col]} vs {[ self.cells[i][col] for i in range(1, col) ]}')
    return self.cells[line][col] not in [ cell_line[col] for cell_line in self.cells[1:line] ]

  def check_3x3(self, line, col):
    line3 = ((line - 1) // 3)
    col3 =  ((col - 1) // 3)

    # print(f'For line {line}, col {col}, checking {self.cells[line][col]} vs {[ self.cells[line3 * 3 + i][col3 * 3 + j] for i in range(1, 4) for j in range(1, 4) if (line3 * 3 + i != line or col3 * 3 + j != col) ]}')

    for i in range(1, 4):
      for j in range(1, 4):
        if (line3 * 3 + i != line or col3 * 3 + j != col) and self.cells[line3 * 3 + i][col3 * 3 + j] == self.cells[line][col]:
          return False

    return True

  def check_black_dot(self, line, col):
    for constraint in self.black_dot_constraints:
      if constraint[0] == (line, col):
        other_line = constraint[1][0]
        other_col = constraint[1][1]

        return self.cells[other_line][other_col] == 0 or self.cells[line][col] * 2 == self.cells[other_line][other_col] or self.cells[line][col] == self.cells[other_line][other_col] * 2

      if constraint[1] == (line, col):
        other_line = constraint[0][0]
        other_col = constraint[0][1]

        return self.cells[other_line][other_col] == 0 or self.cells[line][col] * 2 == self.cells[other_line][other_col] or self.cells[line][col] == self.cells[other_line][other_col] * 2

    return True

  def check_white_dot(self, line, col):
    for constraint in self.white_dot_constraints:
      if constraint[0] == (line, col):
        other_line = constraint[1][0]
        other_col = constraint[1][1]

        return self.cells[other_line][other_col] == 0 or abs(self.cells[line][col] - self.cells[other_line][other_col]) == 1

      if constraint[1] == (line, col):
        other_line = constraint[0][0]
        other_col = constraint[0][1]

        return self.cells[other_line][other_col] == 0 or abs(self.cells[line][col] - self.cells[other_line][other_col]) == 1

    return True

  def check_arrow(self, line, col):
    for constraint in self.arrow_constraints:
      if (line, col) not in constraint:
        continue

      if not strip_zero_unique([ self.cells[cell[0]][cell[1]] for cell in constraint ]):
        return False

      target_sum_point = constraint[0]
      target_sum = self.cells[target_sum_point[0]][target_sum_point[1]]

      found_zero = False
      for cell in constraint[1:]:
        if self.cells[cell[0]][cell[1]] == 0:
          found_zero = True
          break
        target_sum -= self.cells[cell[0]][cell[1]]

      return found_zero or target_sum == 0

    return True

  def check_thermo(self, line, col):
    for constraint in self.thermo_constraints:
      if (line, col) not in constraint:
        continue

      if not strip_zero_unique([ self.cells[cell[0]][cell[1]] for cell in constraint ]):
        return False

      found_zero = False
      for cell in constraint:
        if self.cells[cell[0]][cell[1]] == 0:
          found_zero = True
          break

      thermo_values = [ self.cells[cell[0]][cell[1]] for cell in constraint ]
      return found_zero or all([ thermo_values[i] < thermo_values[i + 1] for i in range(len(thermo_values) - 1) ])

    return True

  def check_palindrome(self, line, col):
    for constraint in self.palindrome_constraints:
      if (line, col) not in constraint:
        continue

      palindrome_values = [ self.cells[cell[0]][cell[1]] for cell in constraint ]

      return all([ palindrome_values[i] == 0 or palindrome_values[i] == palindrome_values[len(palindrome_values) - i - 1] or palindrome_values[len(palindrome_values) - i - 1 == 0] for i in range(len(palindrome_values) // 2) ])

    return True

  def check_box(self, line, col):
    for constraint in self.box_constraints:
      box_sum = constraint[0]

      if (line, col) not in constraint[1]:
        continue

      if not strip_zero_unique([ self.cells[cell[0]][cell[1]] for cell in constraint[1] ]):
        return False

      found_zero = False
      for cell in constraint[1]:
        if self.cells[cell[0]][cell[1]] == 0:
          found_zero = True
          break

      box_values = [ self.cells[cell[0]][cell[1]] for cell in constraint[1] ]
      return found_zero or sum(box_values) == box_sum

    return True

  # check all constraints + sudoku default ones
  def check_at(self, line, col):
    if not self.check_known_points(line, col):
      return False

    if not self.check_line(line, col):
      return False

    if not self.check_col(line, col):
      return False

    if not self.check_3x3(line, col):
      return False

    if not self.check_black_dot(line, col):
      return False

    if not self.check_white_dot(line, col):
      return False

    if not self.check_arrow(line, col):
      return False

    if not self.check_thermo(line, col):
      return False

    if not self.check_box(line, col):
      return False

    if not self.check_palindrome(line, col):
      return False

    return True

  def save_solution(self):
    with open('solution.txt', 'w') as savefile:
      for line in self.cells[1:]:
        savefile.write(' '.join([ str(cell) for cell in line[1:] ]))
        savefile.write('\n')


#encoding: utf-8
'Support subclassing c++ objects in python, with some limitations. Useful primarily for pure-python preprocessors.'

from woo.core import *
from minieigen import *

class PyAttrTrait(object):
	'''
	Class mimicking the `AttrTrait` template in c++, to be used when deriving from :obj:`PyWooObject`, like in this example (can be found in `woo.pyderived` module's source)::

		class _SamplePyDerivedPreprocessor(woo.core.Preprocessor,PyWooObject):
			'Sample preprocessor written in pure python'
			_attrTraits=[
				PyAttrTrait(bool,'switch',True,"Some bool switch, starting group",startGroup='General'),
				PyAttrTrait(int,'a',2,"Integer argument in python"),
				PyAttrTrait(str,'b','whatever',"String argument in python"),
				PyAttrTrait([Node,],'c',[],"List of nodes in python"),
				PyAttrTrait(float,'d',.5,"Parameter with range",range=Vector2(0,.7),startGroup='Advanced'),
				PyAttrTrait(int,'choice',2,"Param with choice",choice=[(0,'choice0'),(1,'choice1'),(2,'choice2')]),
				PyAttrTrait(int,'flags',1,"Param with bits",bits=['bit0','bit1','bit2','bit3','bit4'],buttons=(['Clear flags','self.flags=0','Set flags to zero'],True)),
				PyAttrTrait(int,'choice2',2,'Param with unnamed choice',choice=[0,1,2,3,4,5,6,-1]),
				PyAttrTrait(Vector3,'color',Vector3(.2,.2,.2),"Color parameter",rgbColor=True),
				PyAttrTrait(float,'length',1.5e-3,"Length",unit='mm'),
			]
			def __init__(self,**kw):
				woo.core.Preprocessor.__init__(self)
				self.wooPyInit(self.__class__,woo.core.Preprocessor,**kw)
			def __call__(self):
				pass
				# ...

	This class will be represented like this in the GUI:

	.. image:: fig/pyderived-gui.png

	:param pyType: python type object; can be

		* primitive type (like `bool`, `float`, `Vector3`, `str`, …)
		* sequence of primitive types, written as 1-list: `[bool,]`, `[float,]`, …
		* :obj:`woo Object <woo.core.Object>` or any derived type (:obj:`woo.core.Node`, :obj:`woo.dem.Sphere`, …)
		* sequence of woo objects, e.g. `[woo.core.Node,]`

		When a value is assigned to the attribute, provided type must be convertible to *pyType*, otherwise `TypeError` is raised. There are some additional restrictions:

		* `str` and `unicode` values will *not* be converted to `floats` and such − although python will accept `float('1.23')`

	:param name: name of the attribute, given as `str`
	:param ini: initial (default) value of the attribute
	:param str doc: documentation for this attribute, as it appears in generated docs and tooltips in the UI
	:param str unit: unit given as string, which is looked up in :obj:`woo.unit`; the multiplier is the ratio between *unit* and associated basic unit, which is found automatically.

		.. warning:: `unit` only determines multiplier for the GUI representation; it has no impact on the internal value used; in particular, the *ini* value is *unit-less*. If you want to give units to the initial value, say something like ``PyAttrTrait(float,'angle',60*woo.unit['deg'],units='deg','This is some angle')``.

	:param bool noGui: do not show this attribute in the GUI; use this for attributes which the GUI does not know how to represent (such as python objects, numpy arrays and such), to avoid warnings.
	:param bool noDump: skip this attribute when dumping/loading this object; that means that after loading (or after a :obj:`PyWooObjects.deepcopy`), the attribute will be default-initialized.
	:param bool rgbColor: this attribute is color in the RGB space (its type must be `Vector3`); the GUI will show color picker.
	:param str startGroup: start new attribute group, which is represented as collapsible blocks in the GUI.
	:param str hideIf: python expression which determines whether the GUI hide/show this attribute entry dynamically; use `self` to refer to the instance, as usual.
	:param range: give range (`Vector2` or `Vector2i`) to this numerical (`float` or `int`) attribute − a slider will be shown in the GUI.
	:param choice: this attribute chooses from predefined set of integer values; `choice` itself can be

		* list of unnamed values, e.g. `[0,1,2,3,4]` to choose from;
		* list of named values, e.g. `[(0,'choice0'),(1,'choice1'),(2,'choice2')]`, where the name will be displayed in the GUI, and the number will be assigned when the choice is made.

	:param bits: give names for bit values which this (integer) attribute which represents; they will be shown as array of checkboxes in the GUI.
	:param buttons: Tuple of *list* and *bool*; in the flat list of strings, where each consecutive triplet contains

		1. button label
		2. python expression to evaluate when the button is clicked
		3. label to be shown as button description in the GUI

		The bool at the end determined whether the button is created above (*True*) or below (*False*) the current attribute.

	:param filename: `str` attribute representing filename with file picker in the GUI; the file is possibly non-existent.
	:param existingFilename: `str` attribute for existing filename, with file picker in the GUI.
	:param dirname: `str` attribute for existing directory name, with directory picker in the GUI.
	:param triggerPostLoad: when this attribute is being assigned to, `postLoad(id(self.attr))` will be called.

	'''
	#
	# fake cxxType
	#
	primitiveTypes={int:'int',str:'string',float:'Real',bool:'bool',
		Vector2:'Vector2r',Vector3:'Vector3r',Vector6:'Vector6r',Matrix3:'Matrix3r',Matrix6:'Matrix6r',
		Quaternion:'Quaternionr',Vector2i:'Vector2i',Vector3i:'Vector3i',Vector6i:'Vector6i',
		MatrixX:'MatrixXr',VectorX:'VectorXr'
	}
	def __init__(self,pyType,name,ini,doc,# *,
			unit=None,
			noGui=False,
			noDump=False,
			#multiUnit=False,
			rgbColor=False,
			#prefUnit=None,
			#altUnit=None,
			startGroup=None,
			hideIf=None,
			range=None,
			choice=None,
			bits=None,
			buttons=None,
			altUnits=None,
			filename=False,
			existingFilename=False,
			dirname=False,
			psd=False,
			triggerPostLoad=False,
			guiReadonly=False,
			noGuiResize=False,
		):
		# validity checks
		if range:
			if (pyType,type(range)) not in [(int,Vector2i),(float,Vector2)]: raise TypeError("Range must be Vector2 for floats and Vector2i for ints")
		if isinstance(pyType,list):
			if len(pyType)!=1: raise TypeError('Type must be a list of length exactly one, or a plain type')
			if pyType[0] in self.primitiveTypes: self.cxxType='vector<%s>'%self.primitiveTypes[pyType[0]]
			elif isinstance(pyType[0],type): self.cxxType='vector<shared_ptr<%s>>'%(pyType[0].__name__)
			else: raise TypeError('List element must be a type, not a %s'%(str(pyType[0])))
			# force correct types in the sequence
			if pyType[0] in self.primitiveTypes: ini=[pyType[0](i) for i in ini]
			else:
				for i,v in enumerate(ini):
					if v==None: continue # this is OK
					if not isinstance(v,pyType[0]): raise TypeError("%d-th initial value item must be a %s, not a %s"%(i,pyType[0],type(v)))
		elif pyType in self.primitiveTypes:
			self.cxxType=self.primitiveTypes[pyType]
			ini=pyType(ini)
		elif isinstance(pyType,type):
			self.cxxType='shared_ptr<%s>'%pyType.__name__
		else: raise ValueError('Type must be a list of length one, or a primitive type; not a %s'%str(pyType))
		#
		#
		# mandatory args
		self.pyType=pyType
		self.name=name
		self.ini=ini
		self.doc=doc
		# optional args
		self.noGui=noGui
		self.noDump=noDump
		self.rgbColor=rgbColor
		self.startGroup=startGroup
		self.hideIf=hideIf
		self.range=range
		self.choice=choice
		self.bits=bits
		self.buttons=buttons
		self.filename=filename
		self.existingFilename=existingFilename
		self.dirname=dirname
		self.triggerPostLoad=triggerPostLoad
		self.noGuiResize=noGuiResize
		self.readonly=guiReadonly # this has different meaning in c++ and python, so call it differently
		# those are unsupported in python
		self.noSave=self.hidden=self.pyByRef=self.static=self.activeLabel=self.namedEnum=False
		#
		self.validator=None
		if self.choice and type(self.choice[0])==str:
			#print '%s: Choice of strings (%s)!!!!'%(self.name,str(self.choice))
			# choice from strings
			def validateStrChoice(self,val):
				if val not in self.choice: raise ValueError("%s: '%s' is not an admissible value (must be one of: %s)"%(self.name,str(val),', '.join(["'%s'"%str(c) for c in self.choice])))
			self.validator=validateStrChoice
		# PSD buttons
		if psd:
			if self.buttons: raise ValueError("%s: psd and buttons are mutually exclusive (psd created a button for displaying the PSD)"%self.name)
			self.buttons=(['Plot the PSD','import pylab; pylab.plot(*zip(*self.%s)); pylab.grid(True); pylab.show();'%(self.name),''],0)
		#
		# units
		#
		import woo._units
		baseUnit=woo._units.baseUnit
		def _unicodeUnit(u):
			if isinstance(u,unicode): return u
			elif isinstance(u,str): return unicode(u,'utf-8')
			elif isinstance(u,tuple): return (_unicodeUnit(u[0]),u[1])
			raise ValueError(u"Unknown unit type %sfor %s"%(str(type(u)),u))
		if isinstance(unit,str) or isinstance(unit,unicode):
			unit=[_unicodeUnit(unit)]
			if altUnits: altUnits=[[_unicodeUnit(a) for a in altUnits]]
		if not unit:
			self.unit=None
			self.multiUnit=False
			self.prefUnit=None
			self.altUnits=None
		elif isinstance(unit,list) or isinstance(unit,tuple):
			self.multiUnit=(len(unit)>1)
			self.unit=[]
			self.altUnits=[]
			self.prefUnit=[]
			for u in unit:
				# print self.name,u
				if u not in baseUnit: raise ValueError(u'Unknown unit %s; admissible values are: '%u+', '.join(baseUnit.keys()))
				base=baseUnit[u]
				self.unit.append(base)
				self.prefUnit.append((u,1./woo.unit[u]))
				alts=[]
				for alt in baseUnit:
					if baseUnit[alt]==base and alt!=base: alts.append((alt,1./woo.unit[alt]))
				self.altUnits.append(alts)
			# user-specified alternative units
			if altUnits:
				if len(altUnits)!=len(self.unit): raise ValueError("altUnits must be of the same length as unit")
				for i,au in enumerate(altUnits):
					self.altUnits[i]+=au
		else: raise ValueError('Unknown unit type %s (must be list, tuple, str, unicode): %s'%(type(unit).__name__,str(unit)))
	def validate(self,val):
		'Called when the attribute is set'
		if self.validator: self.validator(self,val)
	def coerceValue(self,val):
		'Check whether *val* has type compatible with declared type (pyType). Raise exception if not. Values converted to required types are returned (it is safe to ignore the return value). In addition, validate the (converted) value, if a validator is defined'
		def tName(T): return (T.__module__+'.' if T.__module__!='__builtin__' else '')+T.__name__
		# sequences
		if isinstance(self.pyType,list):
			assert len(self.pyType)==1
			ret=[]
			if not hasattr(val,'__len__'):
				raise TypeError("Attribute {self.name} declared as sequence of {T}{cxxType}, but its value {val!s} of type {valType} is not a sequence (__len__ not defined).".format(self=self,T=tName(self.pyType[0]),val=val,valType=tName(type(val)),cxxType=((' ('+self.cxxT+' in c++') if hasattr(self,'cxxT') else '')))
			T=self.pyType[0]
			if T in self.primitiveTypes: # check convertibility
				for i,v in enumerate(val):
					try:
						if type(v) in (str,unicode) and T in (float,int): raise TypeError("Don't allow conversions from strings to numbers, since that will fail if used without conversion")
						ret.append(T(v))
					except: raise TypeError("Attribute {self.name} declared as sequence of {T}, but {i}'th item {v!s} of type {itemType} is not convertible to {T}.".format(self=self,i=i,v=v,itemType=tName(type(v)),T=tName(T)))
			else:
				for i,v in enumerate(val):
					ret.append(v)
					if v==None: continue # python representation for NULL shared_ptr
					if not isinstance(v,T): raise TypeError("Attribute {self.name} declared as a sequence of {T}, but {i}'th item {v!s} of type {itemType} is not a {T}.".format(self=self,i=i,v=v,itemType=tName(type(v)),T=tName(T)))
		else:
			# do the same as for sequence items; ugly code duplication
			T=self.pyType
			if T in self.primitiveTypes:
				try:
					if type(val) in (str,unicode) and T in (float,int): raise TypeError("Don't allow conversions from strings to numbers, since that will fail if used without conversion")
					ret=T(val)
				except: raise TypeError("Attribute {self.name} declared as {T}, but value {val!s} of type {valType} is not convertible to {T}".format(self=self,val=val,valType=tName(type(val)),T=tName(T)))
			else:
				# objects
				ret=val
				if val!=None and not isinstance(val,T): raise TypeError("Attribute {self.name} declared as {T}, but value {val!s} of type {valType} is not a {T}".format(self=self,val=val,valType=tName(type(val)),T=tName(T)))
		self.validate(ret)
		return ret
	def __str__(self): return '<PyAttrTrait '+self.name+' @ '+str(id(self))+'>'
	def __repr__(self): return self.__str__()

class PyWooObject:
	'''
	Define some c++-compatibility functions for python classes. Derived class is created as::

		class SomeClass(woo.core.Object,woo.pyderived.PyWooObject): # order of base classes important!
			_attrTraits=[
				# see below
			]
			def __init__(self,**kw):
				woo.core.Object.__init__(self)
				self.wooPyInit(self.__class__,woo.core.Object,**kw)

	This new class automatically obtains several features:

	* dumping/loading via :obj:`woo.core.Object.dump` etc works.
	* the GUI (:obj:`woo.qt.ObjectEditor`) will know how to present this class.
	* documentation for this class will be generated
	* all attributes are type-checked when assigned
	* support for postLoad hooks (see below)

	The `_attrTraits` ist a list of :obj:`PyAttrTrait`; each attribute
	is defined via its traits, which declare its type, default value, documentation and so on -- this
	is documented with :obj:`PyAttrTrait`.

	This example shows trait definitions, and also the `triggerPostLoad` flag::

		class SomeClass(woo.core.Object,woo.pyderived.PyWooObject):
			_PAT=woo.pyderived.PyAttrTrait # alias for class name, to save typing
			_attrTraits=[
				_PAT(float,'aF',1.,'float attr'),
				_PAT([float,],'aFF',[0.,1.,2.],'list of floats attr'),
				_PAT(Vector2,'aV2',(0.,1.),'vector2 attr'),
				_PAT([Vector2,],'aVV2',[(0.,0.),(1.,1.)],'list of vector2 attr'),
				_PAT(woo.core.Node,'aNode',woo.core.Node(pos=(1,1,1)),'node attr'),
				_PAT([woo.core.Node,],'aNNode',[woo.core.Node(pos=(1,1,1)),woo.core.Node(pos=(2,2,2))],'List of nodes'),
				_PAT(float,'aF_trigger',1.,triggerPostLoad=True,doc='Float triggering postLoad'),
			]
			def postLoad(self,I):
				if I==None: pass                  # called when constructed/loaded
				elif I==id(self.aF_trigger): pass # called when aF_trigger is modified
			def __init__(self,**kw):
				pass
				# ...

	The `postLoad` function is called with

	* `None` when the instance has just been created (or loaded); it *shoud* be idempotent, i.e. calling `postLoad(None)` the second time should have no effect::

		SomeClass()         # default-constructed; will call postLoad(None)
		SomeClass(aF=3.)    # default-construct, assign, call postLoad(None)

	* `id(self.attr)` when `self.attr` is modified; this can be used to check for some particular conditions or modify other variables::

		instance=SomeClass()  # calls instance.postLoad(None)
		instance.aF_trigger=3 # calls instance.postLoad(id(instance.aF_trigger))

	  .. note:: Pay attention to not call `postLoad` in infinite regression.

	'''
	def wooPyInit(self,derivedClass,cxxBaseClass,**kw):
		'''Inject methods into derivedClass, so that it behaves like woo.core.Object,
		for the purposes of the GUI and expression dumps'''
		cxxBaseClass.__init__(self) # repeat, just to make sure
		self.cxxBaseClass=cxxBaseClass
		self.derivedClass=derivedClass
		self._instanceTraits={}
		self._attrValues={}
		self._attrTraitsDict=dict([(trait.name,trait) for trait in derivedClass._attrTraits])
		for trait in derivedClass._attrTraits:
			# basic getter/setter
			getter=(lambda self,trait=trait: self._attrValues[trait.name])
			setter=(lambda self,val,trait=trait: self._attrValues.__setitem__(trait.name,trait.coerceValue(val)))
			if trait.triggerPostLoad:
				if not hasattr(derivedClass,'postLoad'): raise RuntimeError('%s.%s declared with triggerPostLoad, but %s.postLoad is not defined.'%(derivedClass.__name__,trait.name,derivedClass.__name__))
				def triggerSetter(self,val,trait=trait):
					self._attrValues[trait.name]=trait.coerceValue(val)
					self.postLoad(id(self._attrValues[trait.name]))
				setter=triggerSetter
			# chain validation and actual setting
			def validatingSetter(self,val,trait=trait,setter=setter):
				trait.validate(val)
				setter(self,val)
			setattr(derivedClass,trait.name,property(getter,validatingSetter,None,trait.doc))
			self._attrValues[trait.name]=trait.ini
		#print derivedClass,self._attrValues
		if kw:
			for k in kw:
				if not hasattr(self,k): raise AttributeError('No such attribute: %s'%k)
				if k in self._attrValues: self._attrValues[k]=self._attrTraitsDict[k].coerceValue(kw[k])
				else: setattr(self,k,kw[k])
		derivedClass.__str__=lambda o:'<%s @ %d (py)>'%(derivedClass.__name__,id(o))
		derivedClass.__repr__=derivedClass.__str__
		derivedClass._cxxAddr=property(lambda sefl: id(self))
		# pickle support
		def __getstate__(self):
			#print '__getstate__ in python'
			ret=self._attrValues.copy()
			ret.update(cxxBaseClass.__getstate__(self)) # get boost::python stuff as well
			return ret
		def __setstate__(self,st):
			#print '__setstate__ in python'
			# set managed attributes indirectly, to avoid side-effects
			for k in st.keys():
				if k in self._attrValues:
					self._attrValues[k]=self._attrTraitsDict[k].coerceValue(st.pop(k))
			# set remaining attributes using setattr
			for k,v in st.items(): setattr(self,k,v)
			# call postLoad as if after loading
			if hasattr(derivedClass,'postLoad'): self.postLoad(None)
		def deepcopy(self):
			'''The c++ dedepcopy uses boost::serialization, we need to use pickle. As long as deepcopy
			is called from python, this function gets precende over the c++ one.'''
			import pickle
			return pickle.loads(pickle.dumps(self))
		derivedClass.__getstate__=__getstate__
		derivedClass.__setstate__=__setstate__
		derivedClass.deepcopy=deepcopy
		if hasattr(derivedClass,'postLoad'):
			self.postLoad(None)

if __name__=='wooMain':
	# do not define this class when running woo normally,
	# so that it does not show up in the preprocessor dialogue

	# inheritance order must not change, due to us using __bases__[0] frequently
	# when traversing the class hierarchy
	class SamplePyDerivedPreprocessor(Preprocessor,PyWooObject):
		'Sample preprocessor written in pure python'
		_classTraits=None
		_attrTraits=[
			PyAttrTrait(bool,'switch',True,"Some bool switch, starting group",startGroup='General'),
			PyAttrTrait(int,'a',2,"Integer argument in python"),
			PyAttrTrait(str,'b','whatever',"String argument in python"),
			PyAttrTrait([Node,],'c',[],"List of nodes in python"),
			PyAttrTrait(float,'d',.5,"Parameter with range",range=Vector2(0,.7),startGroup='Advanced'),
			PyAttrTrait(int,'choice',2,"Param with choice",choice=[(0,'choice0'),(1,'choice1'),(2,'choice2')]),
			PyAttrTrait(int,'flags',1,"Param with bits",bits=['bit0','bit1','bit2','bit3','bit4'],buttons=(['Clear flags','self.flags=0','Set flags to zero'],True)),
			PyAttrTrait(int,'choice2',2,'Param with unnamed choice',choice=[0,1,2,3,4,5,6,-1]),
			PyAttrTrait(Vector3,'color',Vector3(.2,.2,.2),"Color parameter",rgbColor=True),
			PyAttrTrait(float,'length',1.5e-3,"Length",unit='mm'),
			PyAttrTrait(str,'outDir','/tmp',dirname=True,doc='output directory'),
		]
		def __init__(self,**kw):
			# construct all instance attributes
			Preprocessor.__init__(self)
			self.wooPyInit(SamplePyDerivedPreprocessor,Preprocessor,**kw)
		def __call__(self):
			import woo.core, woo.dem, woo.utils
			S=woo.core.Scene(fields=[woo.dem.DemField()])
			S.dem.par.append([
				woo.utils.wall(0,axis=2,sense=1),
				woo.utils.sphere((0,0,1),radius=.2)
			])
			S.dem.gravity=(0,0,-10)
			S.dt=1e-4*woo.utils.pWaveDt(S)
			S.engines=woo.utils.defaultEngines(damping=.01)
			S.dem.collectNodes()
			return S

	import woo.pre
	woo.pre.SamplePyDerivedPreprocessor=SamplePyDerivedPreprocessor

	t=PyAttrTrait(str,'sChoice','aa',choice=['aa','bb','cc'],doc='string choice with validation')
	t.validate('abc')


"""
Artifact module.
"""
from __future__ import annotations

import typing
from typing import Self

from sdk.entities.artifact.metadata import build_metadata
from sdk.entities.artifact.spec import build_spec
from sdk.entities.base.entity import Entity
from sdk.entities.utils.utils import get_uiid
from sdk.utils.api import DTO_ARTF, api_ctx_create, api_ctx_update
from sdk.utils.exceptions import EntityError
from sdk.utils.factories import get_context, get_default_store
from sdk.utils.file_utils import check_file, get_dir
from sdk.utils.uri_utils import get_name_from_uri, get_uri_scheme, rebuild_uri

if typing.TYPE_CHECKING:
    from sdk.entities.artifact.metadata import ArtifactMetadata
    from sdk.entities.artifact.spec import ArtifactSpec

class Artifact(Entity):
    """
    A class representing a artifact.
    """

    def __init__(
        self,
        project: str,
        name: str,
        kind: str = None,
        metadata: ArtifactMetadata = None,
        spec: ArtifactSpec = None,
        local: bool = False,
        embedded: bool = False,
        uuid: str = None,
        **kwargs,
    ) -> None:
        """
        Initialize the Artifact instance.

        Parameters
        ----------
        project : str
            Name of the project.
        name : str
            Name of the artifact.
        kind : str
            Kind of the artifact
        metadata : ArtifactMetadata
            Metadata of the object.
        spec : ArtifactSpec
            Specification of the object.
        local: bool
            If True, run locally.
        embedded: bool
            If True embed object in backend.
        **kwargs
            Keyword arguments.
        """
        super().__init__()
        self.project = project
        self.name = name
        self.kind = kind if kind is not None else "artifact"
        self.metadata = metadata if metadata is not None else build_metadata(name=name)
        self.spec = spec if spec is not None else build_spec(self.kind, **{})
        self.embedded = embedded
        self.id = uuid if uuid is not None else get_uiid()

        self._local = local

        # Temporary local artifact path (see as_file())
        self._temp_path = None

        # Set new attributes
        self._any_setter(**kwargs)

        # Set context
        self._context = get_context(self.project)

        # Set key in spec store://<project>/artifacts/<kind>/<name>:<uuid>
        self.spec.key = (
            f"store://{self.project}/artifacts/{self.kind}/{self.name}:{self.id}"
        )

    #############################
    #  Save / Export
    #############################

    def save(self, uuid: str = None) -> dict:
        """
        Save artifact into backend.

        Parameters
        ----------
        uuid : str
            UUID.

        Returns
        -------
        dict
            Mapping representation of Artifact from backend.
        """
        if self._local:
            raise EntityError("Use .export() for local execution.")

        obj = self.to_dict()

        if uuid is None:
            api = api_ctx_create(self.project, DTO_ARTF)
            return self._context.create_object(obj, api)

        self.id = uuid
        api = api_ctx_update(self.project, DTO_ARTF, self.name, uuid)
        return self._context.update_object(obj, api)

    def export(self, filename: str = None) -> None:
        """
        Export object as a YAML file.

        Parameters
        ----------
        filename : str
            Name of the export YAML file. If not specified, the default value is used.

        Returns
        -------
        None
        """
        obj = self.to_dict()
        filename = (
            filename
            if filename is not None
            else f"artifact_{self.project}_{self.name}.yaml"
        )
        self._export_object(filename, obj)

    #############################
    #  Artifacts Methods
    #############################

    def as_file(self, target: str = None) -> str:
        """
        Get artifact as file. In the case of a local store, the store returns the current
        path of the artifact. In the case of a remote store, the artifact is downloaded in
        a temporary directory.

        Parameters
        ----------
        target : str
            Target path is the remote path of the artifact where it is stored

        Returns
        -------
        str
            Temporary path of the artifact.
        """
        # Get store
        store = get_default_store()

        # If local store, return local artifact path
        if store.is_local():
            self._check_src()
            return self.spec.src_path

        # Check if target path is specified
        self._check_target(target)

        # Check if target path is remote
        self._check_remote()

        # Download artifact and return path
        self._temp_path = store.download(self.spec.target_path)
        return self._temp_path

    def download(
        self, target: str = None, dst: str = None, overwrite: bool = False
    ) -> str:
        """
        Download artifact from backend.

        Parameters
        ----------
        target : str
            Target path is the remote path of the artifact
        dst : str
            Destination path as filename
        overwrite : bool
            Specify if overwrite an existing file

        Returns
        -------
        str
            Path of the downloaded artifact.
        """

        # Check if target path is specified
        self._check_target(target)

        # Check if target path is remote
        self._check_remote()

        # Check if download destination path is specified and rebuild it if necessary
        dst = self._rebuild_dst(dst)

        # Check if destination path exists for overwrite
        self._check_overwrite(dst, overwrite)

        # Get store
        store = get_default_store()

        # Download artifact and return path
        return store.download(self.spec.target_path, dst)

    def upload(self, source: str = None, target: str = None) -> str:
        """
        Upload artifact to backend.

        Parameters
        ----------
        source : str
            Source path is the local path of the artifact
        target : str
            Target path is the remote path of the artifact

        Returns
        -------
        str
            Path of the uploaded artifact.
        """
        # Check if source path is provided.
        self._check_src(source)

        # Check if source path is local
        self._check_local()

        # Check if target path is provided.
        self._check_target(target, upload=True)

        # Check if target path is remote
        self._check_remote()

        # Get store
        store = get_default_store()

        # Upload artifact and return remote path
        return store.upload(self.spec.src_path, self.spec.target_path)

    #############################
    #  Private Helpers
    #############################

    def _check_target(self, target: str = None, upload: bool = False) -> None:
        """
        Check if target path is specified.

        Parameters
        ----------
        target : str
            Target path is the remote path of the artifact

        upload : bool
            Specify if target path is for upload

        Returns
        -------
        None
        """
        if self.spec.target_path is None:
            if target is None:
                if not upload:
                    raise EntityError("Target path is not specified.")
                path = get_dir(self.spec.src_path)
                filename = get_name_from_uri(self.spec.src_path)
                target_path = rebuild_uri(f"{path}/{filename}")
                self.spec.target_path = target_path
                return
            self.spec.target_path = target
            return

    def _check_src(self, src: str = None) -> None:
        """
        Check if source path is specified.

        Parameters
        ----------
        src : str
            Source path is the local path of the artifact

        Returns
        -------
        None

        Raises
        ------
        Exception
            If source path is not specified.
        """
        if self.spec.src_path is None:
            if src is None:
                raise EntityError("Source path is not specified.")
            self.spec.src_path = src

    def _check_remote(self) -> None:
        """
        Check if target path is remote.

        Parameters
        ----------
        ignore_raise : bool
            Specify if raise an exception if target path is not remote

        Returns
        -------
        None

        Raises
        ------
        Exception
            If target path is not remote.
        """
        if self.spec.target_path is None:
            return
        if get_uri_scheme(self.spec.target_path) in ["", "file"]:
            raise EntityError("Only remote source URIs are supported for target paths")

    def _check_local(self) -> None:
        """
        Check if source path is local.

        Returns
        -------
        None

        Raises
        ------
        Exception
            If source path is not local.
        """
        if get_uri_scheme(self.spec.src_path) not in ["", "file"]:
            raise EntityError("Only local paths are supported for source paths.")

    def _rebuild_dst(self, dst: str = None) -> None:
        """
        Check if destination path is specified.

        Parameters
        ----------
        dst : str
            Destination path as filename

        Returns
        -------
        str
            Destination path as filename.
        """
        if dst is None:
            dst = f"./{get_name_from_uri(self.spec.target_path)}"
        return dst

    @staticmethod
    def _check_overwrite(dst: str, overwrite: bool) -> None:
        """
        Check if destination path exists for overwrite.

        Parameters
        ----------
        dst : str
            Destination path as filename.
        overwrite : bool
            Specify if overwrite an existing file.

        Raises
        ------
        Exception
            If destination path exists and overwrite is False.
        """
        if check_file(dst) and not overwrite:
            raise EntityError(f"File {dst} already exists.")

    #############################
    #  Getters and Setters
    #############################

    @property
    def local(self) -> bool:
        """
        Get local flag.
        """
        return self._local

    @property
    def temp_path(self) -> str:
        """
        Get temporary path.
        """
        return self._temp_path

    #############################
    #  Generic Methods
    #############################

    @classmethod
    def from_dict(cls, obj: dict) -> Self:
        """
        Create object instance from a dictionary.

        Parameters
        ----------
        obj : dict
            Dictionary to create object from.

        Returns
        -------
        Self
            Self instance.
        """
        parsed_dict = cls._parse_dict(obj)
        obj_ = cls(**parsed_dict)
        obj_._local = obj_._context.local
        return obj_

    @staticmethod
    def _parse_dict(obj: dict) -> dict:
        """
        Parse dictionary.

        Parameters
        ----------
        obj : dict
            Dictionary to parse.

        Returns
        -------
        dict
            Parsed dictionary.
        """

        # Mandatory fields
        project = obj.get("project")
        name = obj.get("name")
        if project is None or name is None:
            raise EntityError("Project or name are not specified.")

        # Optional fields
        uuid = obj.get("id")
        kind = obj.get("kind")
        embedded = obj.get("embedded")

        # Build metadata and spec
        spec = obj.get("spec")
        spec = spec if spec is not None else {}
        spec = build_spec(kind=kind, **spec)
        metadata = obj.get("metadata", {"name": name})
        metadata = build_metadata(**metadata)

        return {
            "project": project,
            "name": name,
            "kind": kind,
            "uuid": uuid,
            "metadata": metadata,
            "spec": spec,
            "embedded": embedded,
        }

def artifact_from_parameters(
    project: str,
    name: str,
    description: str = "",
    kind: str = "artifact",
    key: str = None,
    src_path: str = None,
    target_path: str = None,
    local: bool = False,
    embedded: bool = False,
    uuid: str = None,
) -> Artifact:
    """
    Create artifact.

    Parameters
    ----------
    project : str
        Name of the project.
    name : str
        Identifier of the artifact.
    description : str
        Description of the artifact.
    kind : str
        The type of the artifact.
    key : str
        Representation of artfact like store://etc..
    src_path : str
        Path to the artifact on local file system or remote storage.
    targeth_path : str
        Destination path of the artifact.
    local : bool
        Flag to determine if object has local execution.
    embedded : bool
        Flag to determine if object must be embedded in project.
    uuid : str
        UUID.

    Returns
    -------
    Artifact
        Artifact object.
    """
    meta = build_metadata(name=name, description=description)
    spec = build_spec(kind, key=key, src_path=src_path, target_path=target_path)
    return Artifact(
        project=project,
        name=name,
        kind=kind,
        metadata=meta,
        spec=spec,
        local=local,
        embedded=embedded,
        uuid=uuid,
    )

def artifact_from_dict(obj: dict) -> Artifact:
    """
    Create artifact from dictionary.

    Parameters
    ----------
    obj : dict
        Dictionary to create artifact from.

    Returns
    -------
    Artifact
        Artifact object.
    """
    return Artifact.from_dict(obj)


import sys
input = sys.stdin.readline

m = int(input())
x = int(1e9)+7

q = 0
for _ in range(m):
    n,s = map(int,input().split())
    nInv = pow(n,-1,x)
    q = (q + (s*nInv)) % x
print(q)


from __future__ import division

import copy
from vistrails.db.versions.v1_0_1.domain import DBVistrail, DBWorkflow, DBLog, \
    DBRegistry, DBGroup, DBTag, DBAnnotation, DBAction, IdScope

def translateVistrail(_vistrail):
    tag_annotations = {}
    notes_annotations = {}
    thumb_annotations = {}
    upgrade_annotations = {}
    prune_annotations = {}

    del_tag_annotations = {}
    del_notes_annotations = {}
    del_thumb_annotations = {}
    del_upgrade_annotations = {}
    del_prune_annotations = {}

    key_lists = {'__tag__': tag_annotations,
                 '__notes__': notes_annotations,
                 '__thumb__': thumb_annotations,
                 '__upgrade__': upgrade_annotations,
                 '__prune__': prune_annotations}

    del_key_lists = {'__tag__': del_tag_annotations,
                     '__notes__': del_notes_annotations,
                     '__thumb__': del_thumb_annotations,
                     '__upgrade__': del_upgrade_annotations,
                     '__prune__': del_prune_annotations}

    _vistrail.update_id_scope()
    id_scope = _vistrail.idScope

    def update_tags(old_obj, translate_dict):
        new_tags = []
        for (id, (_, tag, is_new, is_dirty)) in tag_annotations.iteritems():
            new_tag = DBTag(id=id, name=tag)
            new_tag.is_new = is_new
            new_tag.is_dirty = is_dirty
            new_tags.append(new_tag)
        return new_tags

    def update_prune(old_obj, translate_dict):
        if old_obj.db_id in prune_annotations:
            (_, prune_val, _, _) = prune_annotations[old_obj.db_id]
            if prune_val == str(True):
                return 1
            elif prune_val == str(False):
                return 0
        return None

    def update_annotations(old_obj, translate_dict):
        new_annotations = [DBAnnotation.update_version(a, translate_dict)
                           for a in old_obj.db_annotations]
        if old_obj.db_id in notes_annotations:
            (id, notes, is_new, is_dirty) = notes_annotations[old_obj.db_id]
            ann = DBAnnotation(id=id,
                               key='__notes__',
                               value=notes)
            ann.is_new = is_new
            ann.is_dirty = is_dirty
            new_annotations.append(ann)
        if old_obj.db_id in upgrade_annotations:
            (id, upgrade, is_new, is_dirty) = \
                upgrade_annotations[old_obj.db_id]
            ann =  DBAnnotation(id=id,
                                key='__upgrade__',
                                value=upgrade)
            ann.is_new = is_new
            ann.is_dirty = is_dirty
            new_annotations.append(ann)
        if old_obj.db_id in thumb_annotations:
            (id, thumb, is_new, is_dirty) = thumb_annotations[old_obj.db_id]
            ann = DBAnnotation(id=id,
                               key='__thumb__',
                               value=thumb)
            ann.is_new = is_new
            ann.is_dirty = is_dirty
            new_annotations.append(ann)
        return new_annotations

    def update_actions(old_obj, translate_dict):
        new_actions = []
        for action in old_obj.db_actions:
            if action.db_id in del_notes_annotations:
                (id, notes, is_new, is_dirty) = \
                    del_notes_annotations[action.db_id]
                ann = DBAnnotation(id=id,
                                   key='__notes__',
                                   value=notes)
                ann.is_new = is_new
                ann.is_dirty = is_dirty
                action.db_deleted_annotations.append(ann)
            if action.db_id in del_upgrade_annotations:
                (id, upgrade, is_new, is_dirty) = \
                    del_upgrade_annotations[action.db_id]
                ann = DBAnnotation(id=id,
                                   key='__upgrade__',
                                   value=upgrade)
                ann.is_new = is_new
                ann.is_dirty = is_dirty
                action.db_deleted_annotations.append(ann)
            if action.db_id in del_thumb_annotations:
                (id, thumb, is_new, is_dirty) = \
                    del_thumb_annotations[action.db_id]
                ann = DBAnnotation(id=id,
                                   key='__thumb__',
                                   value=thumb)
                ann.is_new = is_new
                ann.is_dirty = is_dirty
                action.db_deleted_annotations.append(ann)
            new_actions.append(DBAction.update_version(action, translate_dict))
        return new_actions

    def update_workflow(old_obj, translate_dict):
        return DBWorkflow.update_version(old_obj.db_workflow, translate_dict)

    for a in _vistrail.db_actionAnnotations:
        if a.db_key in key_lists:
            key_lists[a.db_key][a.db_action_id] = \
                (a.db_id, a.db_value, a.is_new, a.is_dirty)
    for a in _vistrail.db_deleted_actionAnnotations:
        if a.db_key in del_key_lists:
            del_key_lists[a.db_key][a.db_action_id] = \
                (a.db_id, a.db_value, a.is_new, a.is_dirty)

    translate_dict = {'DBGroup': {'workflow': update_workflow},
                      'DBVistrail': {'tags': update_tags,
                                     'actions': update_actions},
                      'DBAction': {'annotations': update_annotations,
                                   'prune': update_prune}}
    vistrail = DBVistrail.update_version(_vistrail, translate_dict)
    for (id, (_, tag, is_new, is_dirty)) in del_tag_annotations.iteritems():
        new_tag = DBTag(id=id, name=tag)
        new_tag.is_new = is_new
        new_tag.is_dirty = is_dirty
        vistrail.db_deleted_tags.append(new_tag)

    vistrail.db_version = '1.0.1'
    return vistrail

def translateWorkflow(_workflow):
    def update_workflow(old_obj, translate_dict):
        return DBWorkflow.update_version(old_obj.db_workflow, translate_dict)
    translate_dict = {'DBGroup': {'workflow': update_workflow}}
    workflow = DBWorkflow.update_version(_workflow, translate_dict)
    workflow.db_version = '1.0.1'
    return workflow

def translateLog(_log):
    translate_dict = {}
    log = DBLog.update_version(_log, translate_dict)
    log.db_version = '1.0.1'
    return log

def translateRegistry(_registry):
    translate_dict = {}
    registry = DBRegistry.update_version(_registry, translate_dict)
    registry.db_version = '1.0.1'
    return registry


# encoding: utf-8
"""
@author:  liaoxingyu
@contact: sherlockliao01@gmail.com
"""

import logging
import math
from collections import OrderedDict

import torch
from torch import nn
import torch.nn.functional as F
from torch.autograd.variable import Variable
from fastreid.modeling.ops import MetaConv2d, MetaLinear, MetaBNNorm, MetaINNorm, MetaIBNNorm, MetaGate

from fastreid.layers import (
    IBN,
    SELayer,
    Non_local,
    get_norm,
)
from fastreid.utils.checkpoint import get_missing_parameters_message, get_unexpected_parameters_message
from .build import BACKBONE_REGISTRY
from fastreid.utils import comm

K = 4
logger = logging.getLogger(__name__)
model_urls = {
    '18x': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
    '34x': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
    '50x': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
    '101x': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
    'ibn_18x': 'https://github.com/XingangPan/IBN-Net/releases/download/v1.0/resnet18_ibn_a-2f571257.pth',
    'ibn_34x': 'https://github.com/XingangPan/IBN-Net/releases/download/v1.0/resnet34_ibn_a-94bc1577.pth',
    'ibn_50x': 'https://github.com/XingangPan/IBN-Net/releases/download/v1.0/resnet50_ibn_a-d9d0bb7b.pth',
    'ibn_101x': 'https://github.com/XingangPan/IBN-Net/releases/download/v1.0/resnet101_ibn_a-59ea0ac6.pth',
    'se_ibn_101x': 'https://github.com/XingangPan/IBN-Net/releases/download/v1.0/se_resnet101_ibn_a-fabed4e2.pth',
}

def repackage_hidden(h):
    if type(h) == Variable:
        return Variable(h.data)
    else:
        return tuple(repackage_hidden(v) for v in h)

class Sequential_ext(nn.Module):
    """A Sequential container extended to also propagate the gating information
    that is needed in the target rate loss.
    """

    def __init__(self, *args):
        super(Sequential_ext, self).__init__()
        if len(args) == 1 and isinstance(args[0], OrderedDict):
            for key, module in args[0].items():
                self.add_module(key, module)
        else:
            for idx, module in enumerate(args):
                self.add_module(str(idx), module)

    def __getitem__(self, idx):
        if not (-len(self) <= idx < len(self)):
            raise IndexError('index {} is out of range'.format(idx))
        if idx < 0:
            idx += len(self)
        it = iter(self._modules.values())
        for i in range(idx):
            next(it)
        return next(it)

    def __len__(self):
        return len(self._modules)

    def forward(self, input, opt=None):
        for i, module in enumerate(self._modules.values()):
            input = module(input, opt)
        return input

class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self, inplanes, planes, bn_norm, with_ibn=False, with_se=False,
                 stride=1, downsample=None, reduction=16):
        super(BasicBlock, self).__init__()
        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
        if with_ibn:
            self.bn1 = IBN(planes, bn_norm)
        else:
            self.bn1 = get_norm(bn_norm, planes)
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
        self.bn2 = get_norm(bn_norm, planes)
        self.relu = nn.ReLU(inplace=True)
        if with_se:
            self.se = SELayer(planes, reduction)
        else:
            self.se = nn.Identity()
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        identity = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.se(out)

        if self.downsample is not None:
            identity = self.downsample(x)

        out += identity
        out = self.relu(out)

        return out

class MetaSELayer(nn.Module):
    def __init__(self, channel, reduction=16):
        super(MetaSELayer, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.fc1 = MetaLinear(channel, int(channel / reduction), bias=False)
        self.relu = nn.ReLU()
        self.fc2 = MetaLinear(int(channel / reduction), channel, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x, opt=None):
        b, c, _, _ = x.size()
        y = self.avg_pool(x).view(b, c)
        y = self.relu(self.fc1(y, opt))
        y = self.sigmoid(self.fc2(y, opt)).view(b, c, 1, 1)

        return x * y.expand_as(x)

class Bottleneck2(nn.Module):
    expansion = 4*K

    def __init__(self, inplanes, planes, bn_norm, with_ibn=False, with_se=False,
                 stride=1, downsample=None, reduction=16):
        super(Bottleneck2, self).__init__()
        self.conv1 = MetaConv2d(inplanes * K, planes, kernel_size=1, bias=False, groups=K)
        if with_ibn:
            self.bn1 = MetaIBNNorm(planes)
        else:
            self.bn1 = MetaBNNorm(planes)
        self.conv2 = MetaConv2d(planes, planes, kernel_size=3, stride=stride,
                            padding=1, bias=False, groups=K)
        self.bn2 = MetaBNNorm(planes)
        self.conv3 = MetaConv2d(planes, planes * self.expansion, kernel_size=1, bias=False, groups=K)
        self.bn3 = MetaBNNorm(planes * self.expansion)

        self.relu = nn.ReLU(inplace=True)
        if with_se:
            self.se = SELayer(planes * self.expansion, reduction)
        else:
            self.se = nn.Identity()
        self.downsample = downsample
        self.stride = stride

    def forward(self, x, opt=None):

        residual = x

        out = self.conv1(x, opt)
        out = self.bn1(out, opt)
        out = self.relu(out)

        out = self.conv2(out, opt)
        out = self.bn2(out, opt)
        out = self.relu(out)

        out = self.conv3(out, opt)
        out = self.bn3(out, opt)
        out = self.se(out)

        if self.downsample is not None:
            residual = self.downsample(x, opt)

        out += residual
        out = self.relu(out)

        return out

class Bottleneck(nn.Module):
    expansion = 4

    def __init__(self, inplanes, planes, bn_norm, with_ibn=False, with_se=False,
                 stride=1, downsample=None, reduction=16):
        super(Bottleneck, self).__init__()
        if bn_norm == 'IN':
            norm = MetaINNorm
        else:
            norm = MetaBNNorm
        self.conv1 = MetaConv2d(inplanes, planes, kernel_size=1, bias=False)
        if with_ibn:
            self.bn1 = MetaIBNNorm(planes)
        else:
            self.bn1 = norm(planes)
        self.conv2 = MetaConv2d(planes, planes, kernel_size=3, stride=stride,
                            padding=1, bias=False)
        self.bn2 = norm(planes)
        self.conv3 = MetaConv2d(planes, planes * self.expansion, kernel_size=1, bias=False)
        self.bn3 = norm(planes * self.expansion)
        self.relu = nn.ReLU(inplace=True)
        if with_se:
            self.se = SELayer(planes * self.expansion, reduction)
        else:
            self.se = nn.Identity()
        self.downsample = downsample
        self.stride = stride

    def forward(self, x, opt=None):
        residual = x

        out = self.conv1(x, opt)
        out = self.bn1(out, opt)
        out = self.relu(out)

        out = self.conv2(out, opt)
        out = self.bn2(out, opt)
        out = self.relu(out)

        out = self.conv3(out, opt)
        out = self.bn3(out, opt)
        out = self.se(out)

        if self.downsample is not None:
            residual = self.downsample(x, opt)

        out += residual
        out = self.relu(out)

        return out

class Identity(nn.Module):
    def __init__(self):
        super(Identity, self).__init__()

    def forward(self, x):
        return x, None

class HyperRouter(nn.Module):
    def __init__(self, planes):
        super().__init__()
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.planes = planes
        self.fc1 = MetaLinear(planes, planes//16)
        self.fc2 = MetaLinear(planes//16, planes*K)
        self.fc_classifier = MetaLinear(planes*K, 3)
        self.relu = nn.ReLU()
        self.softmax = nn.Softmax(-1)

    def forward(self, x, opt=None):

        x = self.avgpool(x).squeeze(-1).squeeze(-1)

        weight = self.relu(F.normalize(self.fc1(x, opt), 2, -1))
        weight = self.fc2(weight, opt).reshape(-1, self.planes, K)
        domain_cls_logits = self.fc_classifier(weight.reshape(-1, self.planes*K), opt)
        x = self.softmax(torch.einsum('bi,bil->bl', x, weight))

        return x, domain_cls_logits

class ResNet(nn.Module):
    def __init__(self, last_stride, bn_norm, with_ibn, with_se, with_nl, block, layers, non_layers):
        self.inplanes = 64
        super().__init__()
        self.conv1 = MetaConv2d(3, 64, kernel_size=7, stride=2, padding=3,
                               bias=False)
        self.bn1 = MetaBNNorm(64)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True)
        self.layer1 = self._make_layer(block, 64, layers[0]-1, 1, bn_norm, with_ibn, with_se)

        self.adaptor1_base = block(256, 64, 'IN', False, with_se)
        self.adaptor1_sub = Bottleneck2(256, 64, bn_norm, with_ibn, with_se)
        self.router1 = HyperRouter(256)
        self.invariant_norm1 = MetaBNNorm(256)
        self.specific_norm1 = MetaBNNorm(256)
        self.meta_fuse1 = MetaGate(256)
        self.meta_se1 = MetaSELayer(256)
        self.map1 = MetaBNNorm(256, bias_freeze=True)

        self.layer2 = self._make_layer(block, 128, layers[1]-1, 2, bn_norm, with_ibn, with_se)

        self.adaptor2_base = block(512, 128, 'IN', False, with_se)
        self.adaptor2_sub = Bottleneck2(512, 128, bn_norm, with_ibn, with_se)
        self.router2 = HyperRouter(512)
        self.invariant_norm2 = MetaBNNorm(512)
        self.specific_norm2 = MetaBNNorm(512)
        self.meta_fuse2 = MetaGate(512)
        self.meta_se2 = MetaSELayer(512)
        self.map2 = MetaBNNorm(512, bias_freeze=True)

        self.layer3 = self._make_layer(block, 256, layers[2]-1, 2, bn_norm, with_ibn, with_se)

        self.adaptor3_base = block(1024, 256, 'IN', False, with_se)
        self.adaptor3_sub = Bottleneck2(1024, 256, bn_norm, with_ibn, with_se)
        self.router3 = HyperRouter(1024)
        self.invariant_norm3 = MetaBNNorm(1024)
        self.specific_norm3 = MetaBNNorm(1024)
        self.meta_fuse3 = MetaGate(1024)
        self.meta_se3 = MetaSELayer(1024)
        self.map3 = MetaBNNorm(1024, bias_freeze=True)

        self.layer4 = self._make_layer(block, 512, layers[3]-1, last_stride, bn_norm, with_se=with_se)

        self.adaptor4_base = block(2048, 512, 'IN', False, with_se)
        self.adaptor4_sub = Bottleneck2(2048, 512, bn_norm, with_ibn, with_se)
        self.router4 = HyperRouter(2048)
        self.invariant_norm4 = MetaBNNorm(2048)
        self.specific_norm4 = MetaBNNorm(2048)
        self.meta_fuse4 = MetaGate(2048)
        self.meta_se4 = MetaSELayer(2048)
        self.map4 = MetaBNNorm(2048, bias_freeze=True)

        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))

        # Standard Params
        self.softmax = nn.Softmax(1)
        self.sigmoid = nn.Sigmoid()
        self.relu = nn.ReLU()

        self.random_init()

        # fmt: off
        if with_nl: self._build_nonlocal(layers, non_layers, bn_norm)
        else:       self.NL_1_idx = self.NL_2_idx = self.NL_3_idx = self.NL_4_idx = []
        # fmt: on

    def _make_layer(self, block, planes, blocks, stride=1, bn_norm="BN", with_ibn=False, with_se=False):
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = Sequential_ext(
                MetaConv2d(self.inplanes, planes * block.expansion,
                          kernel_size=1, stride=stride, bias=False),
                MetaBNNorm(planes * block.expansion),
            )

        layers = []
        layers.append(block(self.inplanes, planes, bn_norm, with_ibn, with_se, stride, downsample))
        self.inplanes = planes * block.expansion
        for i in range(1, blocks):
            layers.append(block(self.inplanes, planes, bn_norm, with_ibn, with_se))

        return nn.Sequential(*layers)

    def _build_nonlocal(self, layers, non_layers, bn_norm):
        self.NL_1 = nn.ModuleList(
            [Non_local(256, bn_norm) for _ in range(non_layers[0])])
        self.NL_1_idx = sorted([layers[0] - (i + 1) for i in range(non_layers[0])])
        self.NL_2 = nn.ModuleList(
            [Non_local(512, bn_norm) for _ in range(non_layers[1])])
        self.NL_2_idx = sorted([layers[1] - (i + 1) for i in range(non_layers[1])])
        self.NL_3 = nn.ModuleList(
            [Non_local(1024, bn_norm) for _ in range(non_layers[2])])
        self.NL_3_idx = sorted([layers[2] - (i + 1) for i in range(non_layers[2])])
        self.NL_4 = nn.ModuleList(
            [Non_local(2048, bn_norm) for _ in range(non_layers[3])])
        self.NL_4_idx = sorted([layers[3] - (i + 1) for i in range(non_layers[3])])

    def get_all_conv_layers(self, module):
        for m in module:
            if isinstance(m, Bottleneck):
                for _m in m.modules():
                    if isinstance(_m, nn.Conv2d):
                        yield _m

    def forward(self, x, epoch, opt=None):

        x = self.conv1(x, opt)
        x = self.bn1(x, opt)
        x = self.relu(x)
        x = self.maxpool(x)

        weights = []
        out_features = []

        # layer 1
        NL1_counter = 0
        if len(self.NL_1_idx) == 0:
            self.NL_1_idx = [-1]
        for i in range(len(self.layer1)):
            x = self.layer1[i](x, opt)
            if i == self.NL_1_idx[NL1_counter]:
                _, C, H, W = x.shape
                x = self.NL_1[NL1_counter](x)
                NL1_counter += 1

        x_invariant = self.adaptor1_base(x, opt)
        N, C, H, W = x_invariant.shape
        x_specific = self.adaptor1_sub(x.repeat(1, K, 1, 1), opt).reshape(N, K, C, H, W)
        weight, domain_cls_logit = self.router1(x, opt)
        weights.append(weight)
        x_specific = (x_specific * weight.reshape(-1, K, 1, 1, 1)).sum(1)
        x_invariant = self.invariant_norm1(x_invariant, opt)
        x_specific = self.specific_norm1(x_specific, opt)
        x = self.meta_fuse1(x_invariant, x_specific, opt)
        x = self.meta_se1(x, opt)
        temp = self.map1(self.avgpool(x), opt)
        out_features.append(F.normalize(temp, 2, 1)[..., 0, 0])

        # layer 2
        NL2_counter = 0
        if len(self.NL_2_idx) == 0:
            self.NL_2_idx = [-1]
        for i in range(len(self.layer2)):
            x = self.layer2[i](x, opt)
            if i == self.NL_2_idx[NL2_counter]:
                _, C, H, W = x.shape
                x = self.NL_2[NL2_counter](x)
                NL2_counter += 1

        x_invariant = self.adaptor2_base(x, opt)
        N, C, H, W = x_invariant.shape
        x_specific = self.adaptor2_sub(x.repeat(1, K, 1, 1), opt).reshape(N, K, C, H, W)
        weight, domain_cls_logit = self.router2(x, opt)
        weights.append(weight)
        x_specific = (x_specific * weight.reshape(-1, K, 1, 1, 1)).sum(1)
        x_invariant = self.invariant_norm2(x_invariant, opt)
        x_specific = self.specific_norm2(x_specific, opt)
        x = self.meta_fuse2(x_invariant, x_specific, opt)
        x = self.meta_se2(x, opt)
        temp = self.map2(self.avgpool(x), opt)
        out_features.append(F.normalize(temp, 2, 1)[..., 0, 0])

        # layer 3
        NL3_counter = 0
        if len(self.NL_3_idx) == 0:
            self.NL_3_idx = [-1]
        for i in range(len(self.layer3)):
            x = self.layer3[i](x, opt)
            if i == self.NL_3_idx[NL3_counter]:
                _, C, H, W = x.shape
                x = self.NL_3[NL3_counter](x)
                NL3_counter += 1

        x_invariant = self.adaptor3_base(x, opt)
        N, C, H, W = x_invariant.shape
        x_specific = self.adaptor3_sub(x.repeat(1, K, 1, 1), opt).reshape(N, K, C, H, W)
        weight, domain_cls_logit = self.router3(x, opt)
        weights.append(weight)
        x_specific = (x_specific * weight.reshape(-1, K, 1, 1, 1)).sum(1)
        x_invariant = self.invariant_norm3(x_invariant, opt)
        x_specific = self.specific_norm3(x_specific, opt)
        x = self.meta_fuse3(x_invariant, x_specific, opt)
        x = self.meta_se3(x, opt)
        temp = self.map3(self.avgpool(x), opt)
        out_features.append(F.normalize(temp, 2, 1)[..., 0, 0])

        # layer 4
        NL4_counter = 0
        if len(self.NL_4_idx) == 0:
            self.NL_4_idx = [-1]
        for i in range(len(self.layer4)):
            x = self.layer4[i](x, opt)
            if i == self.NL_4_idx[NL4_counter]:
                _, C, H, W = x.shape
                x = self.NL_4[NL4_counter](x)
                NL4_counter += 1

        x_invariant = self.adaptor4_base(x, opt)
        N, C, H, W = x_invariant.shape
        x_specific = self.adaptor4_sub(x.repeat(1, K, 1, 1), opt).reshape(N, K, C, H, W)
        weight, domain_cls_logit = self.router4(x, opt)
        weights.append(weight)
        x_specific = (x_specific * weight.reshape(-1, K, 1, 1, 1)).sum(1)
        x_invariant = self.invariant_norm4(x_invariant, opt)
        x_specific = self.specific_norm4(x_specific, opt)
        x = self.meta_fuse4(x_invariant, x_specific, opt)
        x = self.meta_se4(x, opt)
        temp = self.map4(self.avgpool(x), opt)
        out_features.append(F.normalize(temp, 2, 1)[..., 0, 0])

        weights = torch.cat(weights, -1)

        return x, weights, out_features

    def random_init(self):
        for name, m in self.named_modules():
            if isinstance(m, MetaConv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                nn.init.normal_(m.weight, 0, math.sqrt(2. / n))
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

def init_pretrained_weights(key):
    """Initializes model with pretrained weights.

    Layers that don't match with pretrained layers in name or size are kept unchanged.
    """
    import os
    import errno
    import gdown

    def _get_torch_home():
        ENV_TORCH_HOME = 'TORCH_HOME'
        ENV_XDG_CACHE_HOME = 'XDG_CACHE_HOME'
        DEFAULT_CACHE_DIR = '~/.cache'
        torch_home = os.path.expanduser(
            os.getenv(
                ENV_TORCH_HOME,
                os.path.join(
                    os.getenv(ENV_XDG_CACHE_HOME, DEFAULT_CACHE_DIR), 'torch'
                )
            )
        )
        return torch_home

    torch_home = _get_torch_home()
    model_dir = os.path.join(torch_home, 'checkpoints')
    try:
        os.makedirs(model_dir)
    except OSError as e:
        if e.errno == errno.EEXIST:
            # Directory already exists, ignore.
            pass
        else:
            # Unexpected OSError, re-raise.
            raise

    filename = model_urls[key].split('/')[-1]

    cached_file = os.path.join(model_dir, filename)

    if not os.path.exists(cached_file):
        logger.info(f"Pretrain model don't exist, downloading from {model_urls[key]}")
        if comm.is_main_process():
            gdown.download(model_urls[key], cached_file, quiet=False)

    comm.synchronize()

    logger.info(f"Loading pretrained model from {cached_file}")
    state_dict = torch.load(cached_file, map_location=torch.device('cpu'))
    #CHANGE Reduction Version
    state_dict = torch.load('/home/pengyi/.cache/torch/checkpoints/resnet50_ibn_a-d9d0bb7b.pth', map_location=torch.device('cpu'))

    return state_dict

@BACKBONE_REGISTRY.register()
def build_meta_dynamic_router_resnet_backbone(cfg):
    """
    Create a ResNet instance from config.
    Returns:
        ResNet: a :class:`ResNet` instance.
    """

    # fmt: off
    pretrain      = cfg.MODEL.BACKBONE.PRETRAIN
    pretrain_path = cfg.MODEL.BACKBONE.PRETRAIN_PATH
    last_stride   = cfg.MODEL.BACKBONE.LAST_STRIDE
    bn_norm       = cfg.MODEL.BACKBONE.NORM
    with_ibn      = cfg.MODEL.BACKBONE.WITH_IBN
    with_se       = cfg.MODEL.BACKBONE.WITH_SE
    with_nl       = cfg.MODEL.BACKBONE.WITH_NL
    depth         = cfg.MODEL.BACKBONE.DEPTH
    # fmt: on

    num_blocks_per_stage = {
        '18x': [2, 2, 2, 2],
        '34x': [3, 4, 6, 3],
        '50x': [3, 4, 6, 3],
        '101x': [3, 4, 23, 3],
    }[depth]

    nl_layers_per_stage = {
        '18x': [0, 0, 0, 0],
        '34x': [0, 0, 0, 0],
        '50x': [0, 2, 3, 0],
        '101x': [0, 2, 9, 0]
    }[depth]

    block = {
        '18x': BasicBlock,
        '34x': BasicBlock,
        '50x': Bottleneck,
        '101x': Bottleneck
    }[depth]

    model = ResNet(last_stride, bn_norm, with_ibn, with_se, with_nl, block,
                   num_blocks_per_stage, nl_layers_per_stage)
    if pretrain:
        # Load pretrain path if specifically
        if pretrain_path:
            try:
                state_dict = torch.load(pretrain_path, map_location=torch.device('cpu'))
                logger.info(f"Loading pretrained model from {pretrain_path}")
            except FileNotFoundError as e:
                logger.info(f'{pretrain_path} is not found! Please check this path.')
                raise e
            except KeyError as e:
                logger.info("State dict keys error! Please check the state dict.")
                raise e
        else:
            key = depth
            if with_ibn: key = 'ibn_' + key
            # if with_se:  key = 'se_' + key

            state_dict = init_pretrained_weights(key)

        model_dict = model.state_dict()

        for k in model_dict.keys():
            if k in state_dict:
                v = state_dict[k]
                if model_dict[k].shape == v.shape:
                    model_dict[k] = v
                else:
                    if len(v.shape) == 1:
                        model_dict[k] = v[:model_dict[k].shape[0]]
                    elif len(v.shape) == 2:
                        model_dict[k] = v[:model_dict[k].shape[0], :model_dict[k].shape[1]]
                    elif len(v.shape) == 3:
                        model_dict[k] = v[:model_dict[k].shape[0], :model_dict[k].shape[1], :model_dict[k].shape[2]]
                    elif len(v.shape) == 4:
                        model_dict[k] = v[:model_dict[k].shape[0], :model_dict[k].shape[1], :model_dict[k].shape[2], :model_dict[k].shape[3]]
                    elif len(v.shape) == 5:
                        model_dict[k] = v[:model_dict[k].shape[0], :model_dict[k].shape[1], :model_dict[k].shape[2], :model_dict[k].shape[3], :model_dict[k].shape[4]]
                    else:
                        raise Exception
            else:
                try:
                    if 'adaptor1_base' in k:
                        if model_dict[k].shape == state_dict['layer1.2'+k[13:]].shape:
                            model_dict[k] = state_dict['layer1.2'+k[13:]]
                            print('Done, adaptor', k)
                        else:
                            print('Skip, adaptor', k)
                    elif 'adaptor1_sub' in k:
                        if 'conv3' in k:
                            v = state_dict['layer1.2'+k[12:]]
                            Cout, Cin, H, W = v.shape
                            model_dict[k] = F.avg_pool1d(v.permute(0, 2, 3, 1).reshape(Cout, H*W, Cin), kernel_size=K).reshape(Cout, H, W, -1).permute(0, 3, 1, 2).repeat(K, 1, 1, 1)
                        elif 'bn3' in k:
                            v = state_dict['layer1.2'+k[12:]]
                            model_dict[k] = v.repeat(K)
                        elif model_dict[k].shape == state_dict['layer1.2'+k[12:]].shape:
                            model_dict[k] = state_dict['layer1.2'+k[12:]]
                        else:
                            v = state_dict['layer1.2'+k[12:]]
                            Cout, Cin, H, W = v.shape
                            model_dict[k] = F.avg_pool1d(v.permute(0, 2, 3, 1).reshape(Cout, H*W, Cin), kernel_size=K).reshape(Cout, H, W, -1).permute(0, 3, 1, 2)
                        print('Done, adaptor', k)
                    elif 'adaptor2_base' in k:
                        if model_dict[k].shape == state_dict['layer2.3'+k[13:]].shape:
                            model_dict[k] = state_dict['layer2.3'+k[13:]]
                            print('Done, adaptor', k)
                        else:
                            print('Skip, adaptor', k)
                    elif 'adaptor2_sub' in k:
                        if 'conv3' in k:
                            v = state_dict['layer2.3'+k[12:]]
                            Cout, Cin, H, W = v.shape
                            model_dict[k] = F.avg_pool1d(v.permute(0, 2, 3, 1).reshape(Cout, H*W, Cin), kernel_size=K).reshape(Cout, H, W, -1).permute(0, 3, 1, 2).repeat(K, 1, 1, 1)
                        elif 'bn3' in k:
                            v = state_dict['layer2.3'+k[12:]]
                            model_dict[k] = v.repeat(K)
                        elif model_dict[k].shape == state_dict['layer2.3'+k[12:]].shape:
                            model_dict[k] = state_dict['layer2.3'+k[12:]]
                        else:
                            v = state_dict['layer2.3'+k[12:]]
                            Cout, Cin, H, W = v.shape
                            model_dict[k] = F.avg_pool1d(v.permute(0, 2, 3, 1).reshape(Cout, H*W, Cin), kernel_size=K).reshape(Cout, H, W, -1).permute(0, 3, 1, 2)
                        print('Done, adaptor', k)

                    elif 'adaptor3_base' in k:
                        if model_dict[k].shape == state_dict['layer3.5'+k[13:]].shape:
                            model_dict[k] = state_dict['layer3.5'+k[13:]]
                            print('Done, adaptor', k)
                        else:
                            print('Skip, adaptor', k)
                    elif 'adaptor3_sub' in k:
                        if 'conv3' in k:
                            v = state_dict['layer3.5'+k[12:]]
                            Cout, Cin, H, W = v.shape
                            model_dict[k] = F.avg_pool1d(v.permute(0, 2, 3, 1).reshape(Cout, H*W, Cin), kernel_size=K).reshape(Cout, H, W, -1).permute(0, 3, 1, 2).repeat(K, 1, 1, 1)
                        elif 'bn3' in k:
                            v = state_dict['layer3.5'+k[12:]]
                            model_dict[k] = v.repeat(K)
                        elif model_dict[k].shape == state_dict['layer3.5'+k[12:]].shape:
                            model_dict[k] = state_dict['layer3.5'+k[12:]]
                        else:
                            v = state_dict['layer3.5'+k[12:]]
                            Cout, Cin, H, W = v.shape
                            model_dict[k] = F.avg_pool1d(v.permute(0, 2, 3, 1).reshape(Cout, H*W, Cin), kernel_size=K).reshape(Cout, H, W, -1).permute(0, 3, 1, 2)
                        print('Done, adaptor', k)

                    elif 'adaptor4_base' in k:
                        if model_dict[k].shape == state_dict['layer4.2'+k[13:]].shape:
                            model_dict[k] = state_dict['layer4.2'+k[13:]]
                            print('Done, adaptor', k)
                        else:
                            print('Skip, adaptor', k)
                    elif 'adaptor4_sub' in k:
                        if 'conv3' in k:
                            v = state_dict['layer4.2'+k[12:]]
                            Cout, Cin, H, W = v.shape
                            model_dict[k] = F.avg_pool1d(v.permute(0, 2, 3, 1).reshape(Cout, H*W, Cin), kernel_size=K).reshape(Cout, H, W, -1).permute(0, 3, 1, 2).repeat(K, 1, 1, 1)
                        elif 'bn3' in k:
                            v = state_dict['layer4.2'+k[12:]]
                            model_dict[k] = v.repeat(K)
                        elif 'bn1' in k:
                            if 'IN' in k:
                                model_dict[k] = state_dict['layer4.2.bn1.'+k.split('.')[-1]][:256]
                            else:
                                model_dict[k] = state_dict['layer4.2.bn1.'+k.split('.')[-1]][256:]
                        elif model_dict[k].shape == state_dict['layer4.2'+k[12:]].shape:
                            model_dict[k] = state_dict['layer4.2'+k[12:]]
                        else:
                            v = state_dict['layer4.2'+k[12:]]
                            Cout, Cin, H, W = v.shape
                            model_dict[k] = F.avg_pool1d(v.permute(0, 2, 3, 1).reshape(Cout, H*W, Cin), kernel_size=K).reshape(Cout, H, W, -1).permute(0, 3, 1, 2)
                        print('Done, adaptor', k)

                except Exception:
                    pass
        incompatible = model.load_state_dict(model_dict, strict=False)
        if incompatible.missing_keys:
            logger.info(
                get_missing_parameters_message(incompatible.missing_keys)
            )
        if incompatible.unexpected_keys:
            logger.info(
                get_unexpected_parameters_message(incompatible.unexpected_keys)
            )

    return model


'''
Write a program that takes 3 integers as input and checks whether they can form the sides of a right angled triangle or not. Print YES if they can form a right angled triangle. NO, otherwise.

Input Format:

Single line of input contains three numbers

Output Format:

Print YES or NO

Example:

Input:

5 4 3

Output:

YES

Example:

Input:

10 20 30

Output:

NO
'''

a = [int(x) for x in input().split()[:3]]
a.sort()
if(a[2]*a[2]==a[0]*a[0]+a[1]*a[1]):
    print("Yes")
else:
    print("No")


import pygame
from pygame.surface import *
from pygame.sprite import Sprite
from pygame.sprite import RenderUpdates as SpriteGroup
from pygame.sprite import spritecollide
from pygame.sprite import spritecollideany
from pygame.rect import Rect
from random import *
from config import *
from log import *

screen = None

def createFrames(image):
    fr_width = image.get_height()
    fr_size = fr_width, fr_width
    frames = []
    for frame_no in range(0, image.get_width(), fr_width):
        frame = pygame.Surface(fr_size)
        frame.blit(image, (0,0), ((frame_no,0), fr_size))
        frame.set_colorkey(PUCE)
        frames.append(frame)
    return frames

def initPygame():
    pygame.init()
    global screen

    screen = pygame.display.set_mode(RESOLUTION, FLAGS)

class Widget(Sprite):
    """Use Widget class for better movement tracking

       Widget class inherits from Sprite class.
       Test cases for class Widget

       >>> from widget import *
       >>> import pygame
       >>> s = pygame.surface.Surface((30,30))
       >>> w = Widget(s, (0,0,30,30), (0,0))
       >>> w.rect
       <rect(0, 0, 30, 30)>
       >>> w.update()
       >>> w.rect
       <rect(0, 0, 30, 30)>
       >>> w.getMovement()
       [0, 0]

       >>> w.setX(1)
       >>> w.getX()
       1

       >>> w.setY(4)
       >>> w.getY()
       4

       >>> w.setMovement((3,5))
       >>> w.getMovement()
       (3, 5)

       >>> w.getName()
       (0, 0)

       >>> w.setPosition((5,7))
       >>> w.getPosition()
       (5, 7)
    """

    def __init__(self, image, rect, name=''):
        """Instantiate a widget with a given surface,
           rectangle, and (x,y) movement pair.
        """
        Sprite.__init__(self)
        self.movement = [0, 0]
        self.rect = Rect(rect)
        self.lastRect = self.rect
        self.image = image
        self.name = name
        self.frames = []
        self.frameIndex = 0
        self.frameRate = 1
        self.timeDelay = WIDGETFRAMES
        self.lastUpdate = 0
        self.world = None
        self.undone = False
        self.id = self.rect.top + self.rect.left +\
                  self.rect.width + self.rect.height

    def attachToWorld(self, world):
        self.world = world
        self.id = self.world.curWidgetID
        self.world.curWidgetID += 1

    def startAnimation(self, frames, startIndex, frameRate):
        self.frames = frames
        self.frameIndex = startIndex
        self.frameRate = frameRate
        self.image = self.frames[startIndex]
        self.lastUpdate = self.timeDelay

    def __str__(self):
        return str(self.rect.left) + str(self.rect.top) + str(self.id)

    def setMovement(self, vector):
        """Set movement with a pair"""
        if(self.movement != [0,0]
           and vector == [0,0]):
            self.world.dirtyGroup.add(self)
        self.movement = vector

    def getMovement(self):
        """Return movement as a pair"""
        return self.movement

    def setStop(self):
        """Set movement to 0"""
        self.setMovement([0,0])

    def setY(self, y):
        """Set y-component of movement"""
        self.movement[1] = y

    def setX(self, x):
        """Set x-component of movement"""
        self.movement[0] = x

    def getX(self):
        """Get x-component of movement"""
        return self.movement[0]
    def getY(self):
        """Set y-component of movement"""
        return self.movement[1]

    def setPosition(self, pair):
        """Set x and y coords of Widget"""
        self.rect.topleft = pair

    def getPosition(self):
        """Get x and y coords of Widget"""
        return self.rect.topleft

    def update(self):
        """Move sprite according to its movement vector"""
        # Widget needs to be animated
        if (len(self.frames) > 0):
            if self.lastUpdate <= 0:
                self.frameIndex = (self.frameIndex+1)%(len(self.frames))
                self.image = self.frames[self.frameIndex]
                self.lastUpdate = self.timeDelay
                self.world.dirtyGroup.add(self)
            else:
                self.lastUpdate -= 1

        elif(self.getMovement != [0,0]):
            self.lastRect = Rect(self.rect)
            self.rect.move_ip(self.movement)
            self.world.dirtyGroup.add(self)

    def undoUpdate(self):
        """Widget returns to state prior to last update()"""
        self.rect = self.lastRect

    def getShadow(self):
        shadow = Sprite()
        shadow.rect = self.lastRect.move(0,0)
        return shadow

    def getName(self):
        """Get name of Widget"""
        return self.name

class WorldlessWidget(Widget):
    def update(self):
        """Move sprite according to its movement vector"""
        # Widget needs to be animated
        if (len(self.frames) > 0):
            if self.lastUpdate <= 0:
                self.frameIndex = (self.frameIndex+1)%(len(self.frames))
                self.image = self.frames[self.frameIndex]
                self.lastUpdate = self.timeDelay

        self.lastRect = Rect(self.rect)
        self.rect.move_ip(self.movement)


# Configuration file for the Sphinx documentation builder.
#
# For the full list of built-in configuration values, see the documentation:
# https://www.sphinx-doc.org/en/master/usage/configuration.html

# -- Project information -----------------------------------------------------
# https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information

import os
import sys

project = 'Country Rank'
copyright = '2023, Giovanni Stephens'
author = 'Giovanni Stephens'
release = '1.0.0'

# The master toctree document.
master_doc = "index"

# -- General configuration ---------------------------------------------------
# https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration

extensions = [
    'sphinx.ext.duration',
    'sphinx.ext.doctest',
    'sphinx.ext.autodoc',
    'sphinx.ext.autosummary',
    'sphinx.ext.viewcode',
    'sphinx.ext.napoleon'
]

templates_path = ['_templates']
exclude_patterns = []

sys.path.insert(0, os.path.abspath('../..'))
sys.path.insert(0, os.path.abspath("_themes"))

# -- Options for HTML output -------------------------------------------------
# https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output

html_theme = 'sphinx_rtd_theme'
html_static_path = ['_static']

add_function_parentheses = False
add_module_names = True
html_show_sourcelink = False
html_show_sphinx = False

html_sidebars = {
   '**': ['globaltoc.html', 'sourcelink.html', 'searchbox.html'],
   'using/windows': ['searchbox.html'],
}