import unittest

from absl.testing import absltest, parameterized

import jax
from jax.config import config
import jax.dlpack
import jax.numpy as jnp
from jax import test_util as jtu

config.parse_flags_with_absl()
FLAGS = config.FLAGS

try:
  import torch
  import torch.utils.dlpack
except ImportError:
  torch = None

try:
  import cupy
except ImportError:
  cupy = None

dlpack_dtypes = [jnp.int8, jnp.int16, jnp.int32, jnp.int64,
                   jnp.uint8, jnp.uint16, jnp.uint32, jnp.uint64,
                   jnp.float16, jnp.float32, jnp.float64]
torch_dtypes = [jnp.int8, jnp.int16, jnp.int32, jnp.int64,
                jnp.uint8, jnp.float16, jnp.float32, jnp.float64]

nonempty_nonscalar_array_shapes = [(4,), (3, 4), (2, 3, 4)]
empty_array_shapes = []
empty_array_shapes += [(0,), (0, 4), (3, 0),]
nonempty_nonscalar_array_shapes += [(3, 1), (1, 4), (2, 1, 4)]

nonempty_array_shapes = [()] + nonempty_nonscalar_array_shapes
all_shapes = nonempty_array_shapes + empty_array_shapes

class DLPackTest(jtu.JaxTestCase):
  def setUp(self):
    super(DLPackTest, self).setUp()
    if jtu.device_under_test() == "tpu":
      self.skipTest("DLPack not supported on TPU")

  @parameterized.named_parameters(jtu.cases_from_list(
     {"testcase_name": "_{}".format(
        jtu.format_shape_dtype_string(shape, dtype)),
     "shape": shape, "dtype": dtype}
     for shape in all_shapes
     for dtype in dlpack_dtypes))
  def testJaxRoundTrip(self, shape, dtype):
    rng = jtu.rand_default(self.rng())
    np = rng(shape, dtype)
    x = jnp.array(np)
    dlpack = jax.dlpack.to_dlpack(x)
    y = jax.dlpack.from_dlpack(dlpack)
    self.assertAllClose(np.astype(x.dtype), y)

    self.assertRaisesRegex(RuntimeError,
                           "DLPack tensor may be consumed at most once",
                           lambda: jax.dlpack.from_dlpack(dlpack))

  @parameterized.named_parameters(jtu.cases_from_list(
     {"testcase_name": "_{}".format(
        jtu.format_shape_dtype_string(shape, dtype)),
     "shape": shape, "dtype": dtype}
     for shape in all_shapes
     for dtype in torch_dtypes))
  @unittest.skipIf(not torch, "Test requires PyTorch")
  def testTorchToJax(self, shape, dtype):
    if not FLAGS.jax_enable_x64 and dtype in [jnp.int64, jnp.float64]:
      self.skipTest("x64 types are disabled by jax_enable_x64")
    rng = jtu.rand_default(self.rng())
    np = rng(shape, dtype)
    x = torch.from_numpy(np)
    x = x.cuda() if jtu.device_under_test() == "gpu" else x
    dlpack = torch.utils.dlpack.to_dlpack(x)
    y = jax.dlpack.from_dlpack(dlpack)
    self.assertAllClose(np, y)

  @parameterized.named_parameters(jtu.cases_from_list(
     {"testcase_name": "_{}".format(
        jtu.format_shape_dtype_string(shape, dtype)),
     "shape": shape, "dtype": dtype}
     for shape in all_shapes
     for dtype in torch_dtypes))
  @unittest.skipIf(not torch, "Test requires PyTorch")
  def testJaxToTorch(self, shape, dtype):
    rng = jtu.rand_default(self.rng())
    np = rng(shape, dtype)
    x = jnp.array(np)
    dlpack = jax.dlpack.to_dlpack(x)
    y = torch.utils.dlpack.from_dlpack(dlpack)
    self.assertAllClose(np, y.numpy())

class CudaArrayInterfaceTest(jtu.JaxTestCase):

  def setUp(self):
    super(CudaArrayInterfaceTest, self).setUp()
    if jtu.device_under_test() != "gpu":
      self.skipTest("__cuda_array_interface__ is only supported on GPU")

  @parameterized.named_parameters(jtu.cases_from_list(
     {"testcase_name": "_{}".format(
        jtu.format_shape_dtype_string(shape, dtype)),
     "shape": shape, "dtype": dtype}
     for shape in all_shapes
     for dtype in dlpack_dtypes))
  @unittest.skipIf(not cupy, "Test requires CuPy")
  def testJaxToCuPy(self, shape, dtype):
    rng = jtu.rand_default(self.rng())
    x = rng(shape, dtype)
    y = jnp.array(x)
    z = cupy.asarray(y)
    self.assertEqual(y.__cuda_array_interface__["data"][0],
                     z.__cuda_array_interface__["data"][0])
    self.assertAllClose(x, cupy.asnumpy(z))

if __name__ == "__main__":
  absltest.main(testLoader=jtu.JaxTestLoader())


from setuptools import setup
import os
import codecs

def fpath(name):
    return os.path.join(os.path.dirname(__file__), name)

def read(fname):
    return codecs.open(fpath(fname), encoding='utf-8').read()

requirements = read(fpath('requirements.txt'))

setup(name='scikit-dsp-comm',
      version='0.0.4',
      description='DSP and Comm package.',
      long_description=read(fpath('README.md')),
      author='Mark Wickert',
      author_email='mwickert@uccs.edu',
      url='https://github.com/mwickert/scikit-dsp-comm',
      package_dir={'sk_dsp_comm': 'sk_dsp_comm'},
      packages=['sk_dsp_comm'],
      package_data={'sk_dsp_comm': ['ca1thru37.txt']},
      license='BSD',
      install_requires=requirements.split(),
      test_suite='nose.collector',
      tests_require=['nose','numpy', 'tox'],
      extras_require={
            'helpers': ['pyaudio', 'pyrtlsdr']
      }
     )


import pygame as pg
import file_paths as fp

screen_width=1600
screen_height=900

class MainMenu:
    def __init__(self,screen):
        self.BtnEffect=pg.mixer.Sound(fp.btneffect)
        self.running=True
        self.screen=screen
        self.TitleScreenTexture=pg.image.load(fp.TitleScreenPath).convert_alpha()
        self.StartBtnTexture=pg.image.load(fp.StartBtnPath).convert_alpha()
        self.HelpBtnTexture=pg.image.load(fp.HelpBtnPath).convert_alpha()
        self.Startrect=self.StartBtnTexture.get_rect()
        self.Startrect.center=(screen_width/2+screen_width/4-50,screen_height/2-20)
        self.Helprect=self.HelpBtnTexture.get_rect()
        self.Helprect.center=(screen_width/2+screen_width/4-50,screen_height/2+self.Helprect.h)
        self.return_val=0
    def run(self):
        while self.running:
            pg.display.update()
            self.events()
            self.draw(self.screen)
    def events(self):
            for event in pg.event.get():
                if event.type==pg.QUIT:
                    self.running = False
                if event.type==pg.KEYDOWN:
                    if event.key==pg.K_SPACE:
                        self.running = False
                        self.return_val=1
                if event.type==pg.MOUSEBUTTONDOWN:
                    if event.button==1:
                        if self.Startrect.collidepoint(event.pos):
                            self.BtnEffect.play()
                            self.running = False
                            self.return_val=1
                        if self.Helprect.collidepoint(event.pos):
                            self.BtnEffect.play()
                            self.running = False
                            self.return_val=2

    def draw(self,screen):
        screen.blit(self.TitleScreenTexture,(0,0))
        screen.blit(self.StartBtnTexture,self.Startrect)
        screen.blit(self.HelpBtnTexture,self.Helprect)


import os
from string import Template
# SETUP
nbdir = "./notebooks/"
rstdir = "./notebooks_rst/"
template = Template(open("notebooks_index.template").read())
exclude_prefixes = ('__', '.')  # exclusion prefixes
file_suffixes = (".ipynb",)
title_levels = "+=~-_"

# DATA PROCESSING (whole section needs serious cleaning)
for dirpath, dirnames, filenames in os.walk(nbdir):
  dirnames[:] = [dirname
                 for dirname in dirnames
                 if not dirname.startswith(exclude_prefixes)]
  filenames = [f for f in filenames if f.endswith(file_suffixes)]
  print("#OPERATING ON: ",dirpath, dirnames, filenames)
  """
  if len(dirnames) + len(filenames)==0:
    print("=> No files to proceed, skipping")
  else:
  """
  if True:
    path_depth = len(dirpath.strip("./").strip("/").split("/")) # Depth of the subdir
    rst_path = dirpath.replace(nbdir, rstdir) # Output rst path
    if rst_path.endswith("/") == False: rst_path += "/"
    if os.path.isdir(rst_path) == False: os.mkdir(rst_path)
    print("rst_path={0}  (path_depth = {1})".format(rst_path, path_depth))
    node = dirpath.strip("/").split("/")[-1]
    node_index = template.substitute(title = node,
                                     underline = 80*title_levels[path_depth])
    for d in dirnames:
      node_index += "   " + d + "/" + d + "\n"
    for f in filenames:
      nb = f[:-6]
      rst_back_path = "/".join([".."]*len(dirpath.strip(
                      "./").strip("/").split("/"))) + "/"
      nb_rst_path = ( rst_path.strip("/").strip("./")) +"/"+ nb +".rst"
      node_index += "   " + nb + "\n"
      os.system("jupyter-nbconvert {0}/{1}.ipynb --to rst --output {2}".format(
            dirpath, nb, rst_back_path+nb_rst_path))
      rst = ""
      rst += ".. Note::\n\n  This notebook can be downloaded here: "
      rst += ":download:`{2}.ipynb <{0}{1}/{2}.ipynb>` \n\n".format(
               rst_back_path, dirpath.strip("./"), nb)
      rst += ".. contents::\n   :depth: 2\n"
      rst += open(nb_rst_path).read()
      open(nb_rst_path, "w").write(rst)
    print("####", node)
    open(rst_path + node + ".rst", "w").write(node_index)


class Cart:
    goods = []

    def add(self, gd):
        self.goods.append(gd)

    def remove(self, indx):
        self.goods.pop(indx)

    def get_list(self):
        list = []
        for i in self.goods:
            strok = i.name + ': ' + str(i.price)
            list.append(strok)
        return list

class Table:

    def __init__(self, name, price):
        self.name = name
        self.price = price

class TV:

    def __init__(self, name, price):
        self.name = name
        self.price = price

class Notebook:

    def __init__(self, name, price):
        self.name = name
        self.price = price

class Cup:

    def __init__(self, name, price):
        self.name = name
        self.price = price

cart = Cart()
tv1 = TV("Samsung", 22800)
tv2 = TV("LG", 14880)
table = Table("IKEA", 20000)
N1 = Notebook("Acer", 30000)
N2 = Notebook("Asus", 40000)
cup = Cup("Cup", 1000)
cart.add(tv1)
cart.add(tv2)
cart.add(table)
cart.add(N1)
cart.add(N2)
cart.add(cup)
print(cart.get_list())


from django.test import TestCase, Client
from django.urls import reverse

class TestViews(TestCase):

    def setUp(self):
        self.client = Client()
        self.driver_list_url = reverse("driver_list")
        self.constructor_list_url = reverse("constructor_list")
        self.circuit_list_url = reverse("circuit_list")

    def test_driver_list_GET(self):
        response = self.client.get(self.driver_list_url)
        self.assertEquals(response.status_code, 200)
        self.assertTemplateUsed(response, "f1app/driver_list.html")

    def test_constructor_list_GET(self):
        response = self.client.get(self.constructor_list_url)
        self.assertEquals(response.status_code, 200)
        self.assertTemplateUsed(response, "f1app/constructor_list.html")

    def test_circuit_list_GET(self):
        response = self.client.get(self.circuit_list_url)
        self.assertEquals(response.status_code, 200)
        self.assertTemplateUsed(response, "f1app/circuit_list.html")


#!/usr/bin/env python
# Distributed under the terms of the MIT License.

"""Module for system classes with assigned terms.

Author: Andrew Tarzia

"""

import pathlib
from dataclasses import dataclass

import stk
from openmm import OpenMMException, app, openmm

from .beads import CgBead, get_cgbead_from_element
from .errors import ForceFieldUnavailableError, ForceFieldUnitError
from .martini import MartiniTopology, get_martini_mass_by_type
from .utilities import (
    cosine_periodic_angle_force,
    custom_excluded_volume_force,
)

class ForcedSystem:
    molecule: stk.Molecule
    force_field_terms: dict[str, tuple]
    vdw_bond_cutoff: int

    def _available_forces(self, force_type: str) -> openmm.Force:
        available = {
            "HarmonicBondForce": openmm.HarmonicBondForce(),
            "HarmonicAngleForce": openmm.HarmonicAngleForce(),
            "PeriodicTorsionForce": openmm.PeriodicTorsionForce(),
            "custom-excl-vol": custom_excluded_volume_force(),
            "CosinePeriodicAngleForce": cosine_periodic_angle_force(),
        }
        if force_type not in available:
            msg = f"{force_type} not in {available.keys()}"
            raise ForceFieldUnavailableError(msg)
        return available[force_type]

    def _add_bonds(self, system: openmm.System) -> openmm.System:
        forces = self.force_field_terms["bond"]
        force_types = {i.force for i in forces}
        for force_type in force_types:
            if "Martini" in force_type:
                continue
            force_function = self._available_forces(force_type)
            system.addForce(force_function)
            for assigned_force in forces:
                if assigned_force.force != force_type:
                    continue
                try:
                    force_function.addBond(
                        particle1=assigned_force.atom_ids[0],
                        particle2=assigned_force.atom_ids[1],
                        length=assigned_force.bond_r.value_in_unit(
                            openmm.unit.nanometer
                        ),
                        k=assigned_force.bond_k.value_in_unit(
                            openmm.unit.kilojoule
                            / openmm.unit.mole
                            / openmm.unit.nanometer**2
                        ),
                    )
                except AttributeError:
                    msg = f"{assigned_force} in bonds does not have units"
                    raise ForceFieldUnitError(msg)

        return system

    def _add_angles(self, system: openmm.System) -> openmm.System:
        forces = self.force_field_terms["angle"]
        force_types = {i.force for i in forces}
        for force_type in force_types:
            if "Martini" in force_type:
                continue
            force_function = self._available_forces(force_type)
            system.addForce(force_function)
            for assigned_force in forces:
                if assigned_force.force != force_type:
                    continue
                try:
                    if force_type == "CosinePeriodicAngleForce":
                        force_function.addAngle(
                            assigned_force.atom_ids[0],
                            assigned_force.atom_ids[1],
                            assigned_force.atom_ids[2],
                            # Order is important here!
                            [
                                assigned_force.angle_k.value_in_unit(
                                    openmm.unit.kilojoule / openmm.unit.mole
                                ),
                                assigned_force.n,
                                assigned_force.b,
                                (-1) ** assigned_force.n,
                            ],
                        )
                    elif force_type == "HarmonicAngleForce":
                        force_function.addAngle(
                            particle1=assigned_force.atom_ids[0],
                            particle2=assigned_force.atom_ids[1],
                            particle3=assigned_force.atom_ids[2],
                            angle=assigned_force.angle.value_in_unit(
                                openmm.unit.radian
                            ),
                            k=assigned_force.angle_k.value_in_unit(
                                openmm.unit.kilojoule
                                / openmm.unit.mole
                                / openmm.unit.radian**2
                            ),
                        )
                except AttributeError:
                    msg = f"{assigned_force} in angles does not have units"
                    raise ForceFieldUnitError(msg)

        return system

    def _add_torsions(self, system: openmm.System) -> openmm.System:
        forces = self.force_field_terms["torsion"]
        force_types = {i.force for i in forces}
        for force_type in force_types:
            if "Martini" in force_type:
                continue
            force_function = self._available_forces(force_type)
            system.addForce(force_function)
            for assigned_force in forces:
                if assigned_force.force != force_type:
                    continue
                try:
                    force_function.addTorsion(
                        particle1=assigned_force.atom_ids[0],
                        particle2=assigned_force.atom_ids[1],
                        particle3=assigned_force.atom_ids[2],
                        particle4=assigned_force.atom_ids[3],
                        periodicity=assigned_force.torsion_n,
                        phase=assigned_force.phi0.value_in_unit(
                            openmm.unit.radian
                        ),
                        k=assigned_force.torsion_k.value_in_unit(
                            openmm.unit.kilojoules_per_mole
                        ),
                    )
                except AttributeError:
                    msg = f"{assigned_force} in torsions does not have units"
                    raise ForceFieldUnitError(msg)

        return system

    def _add_nonbondeds(self, system: openmm.System) -> openmm.System:
        exclusion_bonds = [
            (i.get_atom1().get_id(), i.get_atom2().get_id())
            for i in self.molecule.get_bonds()
        ]
        forces = self.force_field_terms["nonbonded"]
        force_types = {i.force for i in forces}
        for force_type in force_types:
            force_function = self._available_forces(force_type)
            system.addForce(force_function)
            for assigned_force in forces:
                if assigned_force.force != force_type:
                    continue
                try:
                    force_function.addParticle(
                        [
                            assigned_force.sigma.value_in_unit(
                                openmm.unit.nanometer
                            ),
                            assigned_force.epsilon.value_in_unit(
                                openmm.unit.kilojoules_per_mole
                            ),
                        ],
                    )

                except AttributeError:
                    msg = f"{assigned_force} in nonbondeds does not have units"
                    raise ForceFieldUnitError(msg)

            try:
                # This method MUST be after terms are assigned.
                force_function.createExclusionsFromBonds(
                    exclusion_bonds,
                    self.vdw_bond_cutoff,
                )
            except OpenMMException:
                msg = f"{force_type} is missing a definition for a particle."
                raise ForceFieldUnitError(msg)

        return system

    def _add_forces(self, system: openmm.System) -> openmm.System:
        system = self._add_bonds(system)
        system = self._add_angles(system)
        system = self._add_torsions(system)
        return self._add_nonbondeds(system)

    def get_openmm_topology(self) -> app.topology.Topology:
        raise NotImplementedError

    def get_openmm_system(self) -> openmm.System:
        raise NotImplementedError

@dataclass(frozen=True)
class AssignedSystem(ForcedSystem):
    molecule: stk.Molecule
    force_field_terms: dict[str, tuple]
    system_xml: pathlib.Path
    topology_xml: pathlib.Path
    bead_set: dict[str, CgBead]
    vdw_bond_cutoff: int
    mass: float = 10

    def _get_topology_xml_string(self, molecule: stk.Molecule) -> str:
        ff_str = "<ForceField>\n\n"

        at_str = " <AtomTypes>\n"
        re_str = " <Residues>\n"
        re_str += '  <Residue name="ALL">\n'

        present_beads = {}
        for atom in molecule.get_atoms():
            aestring = atom.__class__.__name__
            aid = atom.get_id()
            acgbead = get_cgbead_from_element(
                estring=aestring,
                bead_set=self.bead_set,
            )
            atype = acgbead.bead_type
            aclass = acgbead.bead_class

            if atype not in present_beads:
                present_beads[atype] = acgbead
                at_str += (
                    f'  <Type name="{atype}" '
                    f'class="{aclass}" element="{aestring}" '
                    f'mass="{self.mass}"/>\n'
                )

            re_str += f'   <Atom name="{aid}" type="{atype}"/>\n'

        for bond in molecule.get_bonds():
            a1id = bond.get_atom1().get_id()
            a2id = bond.get_atom2().get_id()

            re_str += f'   <Bond atomName1="{a1id}" atomName2="{a2id}"/>\n'

        at_str += " </AtomTypes>\n\n"
        re_str += "  </Residue>\n"
        re_str += " </Residues>\n\n"

        ff_str += at_str
        ff_str += re_str

        ff_str += "</ForceField>\n"
        return ff_str

    def _write_topology_xml(self, molecule: stk.Molecule) -> None:
        ff_str = self._get_topology_xml_string(molecule)

        with open(self.topology_xml, "w") as f:
            f.write(ff_str)

    def get_openmm_topology(self) -> app.topology.Topology:
        topology = app.topology.Topology()
        chain = topology.addChain()
        residue = topology.addResidue(name="ALL", chain=chain)

        atoms_added = {}
        for atom in self.molecule.get_atoms():
            a_id = atom.get_id()
            a_estring = atom.__class__.__name__
            a_element = app.element.Element.getByAtomicNumber(
                atom.get_atomic_number()
            )
            a_cgbead = get_cgbead_from_element(
                estring=a_estring,
                bead_set=self.bead_set,
            )

            omm_atom = topology.addAtom(
                name=a_cgbead.bead_type,
                element=a_element,
                residue=residue,
                id=str(a_id),
            )
            atoms_added[a_id] = omm_atom

        for bond in self.molecule.get_bonds():
            a1_id = bond.get_atom1().get_id()
            a2_id = bond.get_atom2().get_id()

            topology.addBond(
                atom1=atoms_added[a1_id],
                atom2=atoms_added[a2_id],
            )

        return topology

    def get_openmm_system(self) -> openmm.System:
        self._write_topology_xml(self.molecule)
        forcefield = app.ForceField(self.topology_xml)
        topology = self.get_openmm_topology()
        system = forcefield.createSystem(topology)
        system = self._add_forces(system)

        with open(self.system_xml, "w") as f:
            f.write(openmm.XmlSerializer.serialize(system))

        return system

@dataclass(frozen=True)
class MartiniSystem(ForcedSystem):
    """
    Assign a system using martini_openmm.

    """

    molecule: stk.Molecule
    force_field_terms: dict[str, tuple]
    system_xml: pathlib.Path
    topology_itp: pathlib.Path
    vdw_bond_cutoff: int
    bead_set: dict[str, CgBead]

    def _get_atoms_string(
        self,
        molecule: stk.Molecule,
        molecule_name: str,
    ) -> str:
        atoms_string = (
            "[atoms]\n"
            ";nr  type  resnr  resid  atom  cgnr  charge  mass  total_charge\n"
        )
        for atom in molecule.get_atoms():
            a_estring = atom.__class__.__name__

            a_cgbead = get_cgbead_from_element(
                estring=a_estring,
                bead_set=self.bead_set,
            )
            nr = atom.get_id() + 1
            type_ = a_cgbead.bead_type
            resnr = 1
            resid = molecule_name[:4].upper()
            charge = 0
            total_charge = 0
            # Charge group, set as different for all for now, like in PYRI.
            cgnr = atom.get_id() + 1
            mass = get_martini_mass_by_type(type_)
            atoms_string += (
                f"{nr} {type_} {resnr} {resid} {a_estring} {cgnr} {charge} "
                f"{mass} {total_charge}\n"
            )
        atoms_string += "\n"
        return atoms_string

    def _get_bonds_string(self, molecule: stk.Molecule) -> str:
        string = "[bonds]\n; i j   funct   length\n"
        forces = self.force_field_terms["bond"]
        for assigned_force in forces:
            force_type = assigned_force.force
            if force_type != "MartiniDefinedBond":
                continue
            length = assigned_force.bond_r.value_in_unit(openmm.unit.nanometer)
            k = assigned_force.bond_k.value_in_unit(
                openmm.unit.kilojoule
                / openmm.unit.mole
                / openmm.unit.nanometer**2
            )
            string += (
                f"  {assigned_force.atom_ids[0]+1} "
                f"{assigned_force.atom_ids[1]+1} "
                f"{assigned_force.funct} "
                f"{length} "
                f"{k}\n"
            )
        string += "\n"
        return string

    def _get_angles_string(self, molecule: stk.Molecule) -> str:
        string = "[angles]\n; i j k    funct  ref.angle   force_k\n"
        forces = self.force_field_terms["angle"]

        for assigned_force in forces:
            force_type = assigned_force.force
            if force_type != "MartiniDefinedAngle":
                continue

            angle = assigned_force.angle.value_in_unit(openmm.unit.degrees)
            k = assigned_force.angle_k.value_in_unit(
                openmm.unit.kilojoule
                / openmm.unit.mole
                / openmm.unit.radian**2
            )
            string += (
                f"  {assigned_force.atom_ids[0]+1} "
                f"{assigned_force.atom_ids[1]+1} "
                f"{assigned_force.atom_ids[2]+1} "
                f"{assigned_force.funct} "
                f"{angle} "
                f"{k}\n"
            )

        string += "\n"
        return string

    def _get_torsions_string(self, molecule: stk.Molecule) -> str:
        string = "[dihedrals]\n; i j k l  funct  ref.angle   force_k\n"
        forces = self.force_field_terms["torsion"]
        force_types = {i.force for i in forces}

        for force_type in force_types:
            if force_type != "MartiniDefinedTorsion":
                continue
            print(force_type)
            raise SystemExit()
        string += "\n"
        return string

    def _get_contraints_string(self, molecule: stk.Molecule) -> str:
        string = "[constraints]\n; i j   funct   length\n"
        for constraint in self.force_field_terms["constraints"]:
            string += (
                f"  {constraint[0]} {constraint[1]} {constraint[2]} "
                f"{constraint[3]} {constraint[4]}"
            )
        string += "\n"
        return string

    def _get_exclusions_string(self, molecule: stk.Molecule) -> str:
        string = "[exclusions]\n; i j   funct   length\n"
        for constraint in self.force_field_terms["constraints"]:
            string += (
                f"  {constraint[0]} {constraint[1]} {constraint[2]} "
                f"{constraint[3]} {constraint[4]}"
            )
        string += "\n"
        return string

    def _write_topology_itp(self, molecule: stk.Molecule) -> None:
        molecule_name = self.topology_itp.name.replace(".itp", "")
        string = (
            f";;; {molecule_name}\n"
            "[moleculetype]\n"
            "; Name nrexcl\n"
            f"{molecule_name} {self.vdw_bond_cutoff}\n\n"
        )

        atoms_string = self._get_atoms_string(molecule, molecule_name)
        bonds_string = self._get_bonds_string(molecule)
        angles_string = self._get_angles_string(molecule)
        torsions_string = self._get_torsions_string(molecule)
        constraints_string = self._get_contraints_string(molecule)

        string += atoms_string
        string += bonds_string
        string += angles_string
        string += torsions_string
        string += constraints_string

        with open(self.topology_itp, "w") as f:
            f.write(string)

    def get_openmm_topology(self) -> app.topology.Topology:
        self._write_topology_itp(self.molecule)
        return MartiniTopology(self.topology_itp).get_openmm_topology()

    def get_openmm_system(self) -> openmm.System:
        self._write_topology_itp(self.molecule)
        topology = MartiniTopology(self.topology_itp)
        system = topology.get_openmm_system()

        system = self._add_forces(system)
        with open(self.system_xml, "w") as f:
            f.write(openmm.XmlSerializer.serialize(system))
        return system


#
# From http://hetland.org/python/distance.py
# Levenshtein distance implementation
# by Magnus Lie Hetland
#
def distance(a,b):
    "Calculates the Levenshtein distance between a and b."
    n, m = len(a), len(b)
    if n > m:
        # Make sure n <= m, to use O(min(n,m)) space
        a,b = b,a
        n,m = m,n

    current = range(n+1)
    for i in range(1,m+1):
        previous, current = current, [i]+[0]*m
        for j in range(1,n+1):
            add, delete = previous[j]+1, current[j-1]+1
            change = previous[j-1]
            if a[j-1] != b[i-1]:
                change = change + 1
            current[j] = min(add, delete, change)

    return current[n]


from rest_framework import serializers
from rest_framework.serializers import ModelSerializer, PrimaryKeyRelatedField
from .models import *

class UsuarioSerializer(ModelSerializer):
    class Meta:
        model = Usuario
        fields = ['email', 'senha']

class PessoaFisicaSerializer(ModelSerializer):
    usuario = UsuarioSerializer(required=False)
    usuario_id = PrimaryKeyRelatedField(queryset=Usuario.objects.all(), required=False, write_only=True)

    class Meta:
        model = PessoaFisica
        fields = '__all__'

    def validate(self, data):
        usuario = data.get("usuario")
        usuario_id = data.get("usuario_id")

        if usuario and usuario_id:
            raise serializers.ValidationError("Você não pode fornecer dados para 'usuario' e 'usuario_id' ao mesmo tempo.")

        if not (usuario or usuario_id):
            raise serializers.ValidationError("Você deve fornecer dados para 'usuario' ou 'usuario_id'.")

        return data

    def create(self, validated_data):
        usuario_data = validated_data.pop('usuario', None)
        usuario_id = validated_data.pop('usuario_id', None)

        if usuario_id:
            usuario = usuario_id
        else:
            usuario = Usuario.objects.create(**usuario_data)

        pessoa_fisica = PessoaFisica.objects.create(usuario = usuario, **validated_data)
        return pessoa_fisica

class ProfissionalSerializer(ModelSerializer):
    class Meta:
        model = Profissional
        fields = '__all__'

class SolicitanteSerializer(ModelSerializer):
    profissional = ProfissionalSerializer(required=False)
    profissional_id = PrimaryKeyRelatedField(queryset=Profissional.objects.all(), required=False, write_only=True)
    class Meta:
        model = Solicitante
        fields = '__all__'

    def validate(self, data):
        profissional = data.get('profissional')
        profissional_id = data.get('profissional_id')

        if profissional and profissional_id:
            raise serializers.ValidationError("Você não pode fornecer dados para 'profissional' e 'profissional_id' ao mesmo tempo.")

        if not (profissional or profissional_id):
            raise serializers.ValidationError("Você deve fornecer dados para 'profissional' ou 'profissional_id'.")

        return data

    def create(self, validated_data):
        profissional_data = validated_data.pop('profissional', None)
        profissional_id   = validated_data.pop('profissional_id', None)

        if profissional_id:
            profissional = profissional_id
        else:
            profissional = Profissional.objects.create(**profissional_data)

        solicitante = Solicitante.objects.create(profissional = profissional, **validated_data)

        return solicitante

class OperacionalSerializer(ModelSerializer):
    profissional = ProfissionalSerializer(required=False)
    profissional_id = PrimaryKeyRelatedField(queryset=Profissional.objects.all(), required=False, write_only=True)

    class Meta:
        model = Operacional
        fields = '__all__'

    def validate(self, data):
        profissional = data.get('profissional')
        profissional_id = data.get('profissional_id')

        if profissional and profissional_id:
            raise serializers.ValidationError("Você não pode fornecer dados para 'profissional' e 'profissional_id' ao mesmo tempo.")

        if not (profissional or profissional_id):
            raise serializers.ValidationError("Você deve fornecer dados para 'profissional' ou 'profissional_id'.")

        return data

    def create(self, validated_data):
        profissional_data = validated_data.pop('profissional', None)
        profissional_id   = validated_data.pop('profissional_id', None)

        if profissional_id:
            profissional = profissional_id
        else:
            profissional = Profissional.objects.create(**profissional_data)

        operacional = Operacional.objects.create(profissional = profissional, **validated_data)

        return operacional

class EnderecoSerializer(ModelSerializer):
    class Meta:
        model = Endereco
        fields = '__all__'
class InstituicaoSerializer(ModelSerializer):
    endereco = EnderecoSerializer()

    class Meta:
        model = Instituicao
        fields = '__all__'

    def create(self, validated_data):
        endereco_data = validated_data.pop('endereco')
        endereco = Endereco.objects.create(**endereco_data)

        instituicao = Instituicao.objects.create(endereco=endereco,**validated_data)
        return instituicao


#!/usr/bin/python
import smbus
import time
# Define some constants from the datasheet
DEVICE     = 0x23 # Default device I2C address
DEVICE2    = 0X5c
POWER_DOWN = 0x00 # No active state
POWER_ON   = 0x01 # Power on
RESET      = 0x07 # Reset data register value
ONE_TIME_HIGH_RES_MODE = 0x20

bus = smbus.SMBus(1)  # Rev 2 Pi uses 1

def convertToNumber(data):
    # Simple function to convert 2 bytes of data
    return ((data[1] + (256 * data[0])) / 1.2)

def readLight(addr=DEVICE):
   # Read data from I2C interface
   data = bus.read_i2c_block_data(addr,ONE_TIME_HIGH_RES_MODE)
   return convertToNumber(data)

def main():

    while True:
        lightLevel=readLight()
        #lightLevel2=readLight(addr=DEVICE2)
        print("Light Level : " + format(lightLevel,'.2f') + " lx")
        #print("Light Level 2 : " + format(lightLevel2,'.2f') + " lx")
        time.sleep(4)

if __name__=="__main__":
    main()


########### below only needed if pydtc not installed. #################
import sys, os
sys.path.insert(0,
    os.path.abspath(os.path.join(os.path.dirname( __file__ ), '..')))
#######################################################################
import time, os
from pydtc import timeout, retry, DTCTimeoutException

@retry(DTCTimeoutException, retries=2, delay=2)
@timeout(4)
def test(err_within_timerange=False):
    for i in range(10):
        print(f'{i+1} seconds.')
        # print(os.getpid())
        time.sleep(1)
        if err_within_timerange:
            raise OSError('exception from function internally.')

if __name__ == "__main__":
    try:
        test()
    except DTCTimeoutException as e:
        print(e)

    try:
        test(True)
    except DTCTimeoutException as e:
        print(e)


import json
import os
import sqlite3
from dataclasses import dataclass
from sqlite3 import Cursor
from tkinter import IntVar

_DB_PATH = "test.db"

@dataclass
class Tag:
    name: str
    is_hidden: bool
    is_selected: IntVar

@dataclass
class DisplayInfo:
    title: str
    info: str

@dataclass
class File:
    name: str
    path: str
    tags: list[str]

    def get_display_info(self) -> tuple[DisplayInfo, ...]:
        tags_info = ""
        for tag in self.tags:
            tags_info += f"{tag}, " if tag != self.tags[-1] else tag
        return DisplayInfo("Name", self.name), DisplayInfo("Path", self.path), DisplayInfo("Tags", tags_info)

def init() -> None:
    if not os.path.exists(_DB_PATH):  # noqa: PTH110
        _create_new()

def debug() -> None:
    print("\nDB DEBUG:")
    print("files:\n", _fetch_files())
    print("\ntags:\n", _fetch_tags())

def fetch_files() -> list[File]:
    files = _fetch_files()
    file_results = []

    for name, path, tags in files:
        file_results.append(File(name, path, json.loads(tags)))

    return file_results

# TAGS

def fetch_tags() -> list[Tag]:
    tags = _fetch_tags()
    content_tags = []

    for tag, is_hidden, is_selected_in_db in tags:
        is_selected = IntVar()
        is_selected.set(is_selected_in_db)

        content_tags.append(Tag(name=tag, is_hidden=is_hidden, is_selected=is_selected))

    return content_tags

def create_tag(tag: str) -> None:
    sql = f"INSERT INTO tags(name, is_hidden, is_selected) VALUES ('{tag}', 0, 0)"  # noqa: S608
    _get_cursor().execute(sql).connection.commit()

def disable_tag_visibility(tag: str) -> None:
    _get_cursor().execute(f"UPDATE tags SET is_hidden=1 WHERE name='{tag}'").connection.commit()  # noqa: S608

def enable_tag_visibility(tag: str) -> None:
    _get_cursor().execute(f"UPDATE tags SET is_hidden=0 WHERE name='{tag}'").connection.commit()  # noqa: S608

def toggle_tag_selection(tag: Tag) -> None:
    _get_cursor().execute(
        f"UPDATE tags SET is_selected={tag.is_selected.get()} WHERE name='{tag.name}'",  # noqa: S608
    ).connection.commit()

def disable_tag_selection(tag: str) -> None:
    _get_cursor().execute(f"UPDATE tags SET is_selected=0 WHERE name='{tag}'").connection.commit()  # noqa: S608

def clear_all_tag_selections() -> None:
    _get_cursor().execute("UPDATE tags SET is_selected=0 ").connection.commit()

# PRIVATES

def _create_new() -> None:
    db_commands = (
        "CREATE TABLE files(name TEXT NOT NULL, path TEXT NOT NULL, tags JSON_DUMPS NOT NULL)",
        "CREATE TABLE tags(name TEXT NOT NULL, is_hidden BOOLEAN NOT NULL, is_selected BOOLEAN NOT NULL)",
        f"""
        INSERT INTO files(name, path, tags)
        VALUES
        (
            'somewhat_really_long_sample_video_name',
            'somewhat_really_long_sample_video_name.mp4',
            '{json.dumps(('Favorite', 'Archive'))}'),
        (
            'non_existent_video',
            'video.dne',
            '{json.dumps(('Archive',))}'
        )
        """,
        "INSERT INTO tags(name, is_hidden, is_selected) VALUES ('Favorite', 0, 0), ('Archive', 0, 0)",
    )

    cursor = _get_cursor()

    for cmd in db_commands:
        cursor.execute(cmd)

    cursor.connection.commit()

def _fetch_tags() -> list[tuple[str, bool, bool]]:
    return _get_cursor().execute("SELECT * FROM tags").fetchall()

def _fetch_files() -> list[tuple[str, str, str]]:
    return _get_cursor().execute("SELECT * FROM files").fetchall()

def _get_cursor() -> Cursor:
    return sqlite3.connect(_DB_PATH).cursor()


import streamlit as st

from Session.StatementConfig import StatementConstants
from Session.UpdateDbSession import UpdateDbSession
from library_hvac_app.list_custom_functions import to_list

class AddSQLTableView:
	def __init__(self, key):
		col1, col2, col3, col4 = st.columns(4)
		self.key = key

		with col1:
			self.create_view_db_button = st.button("Create View DB", key=f"create_view_db button {key}")
			self.create_table_db_button = st.button("Create Table DB", key=f"create_table_db button {key}")
			st.button("Update db", key=f"update_table_db button {key}", on_click=UpdateDbSession.Update_sql_sessionData)

		with col2:
			self.new_view_table_name = st.text_input("Enter New View Table Name", key=f"new_view_table_name {key}")
			self.new_view_comments = st.text_input("Enter Comments to view", key=f"new_view_comments {key}")

		with col3:
			self.category_name = self.__chose_category_type()

		with col4:
			self._add_category()

	def _add_category(self):
		category_name = st.text_input("Type new category name", value="Heat Balance", key=f"category name {self.key}")
		add_category_button = st.button("Add Category", key=f"add category button{self.key}")
		if add_category_button:
			category_view_list = st.session_state[StatementConstants.category_view_list]
			if category_name not in category_view_list:
				category_view_list.append(category_name)

	def __chose_category_type(self):
		category = to_list(st.session_state[StatementConstants.category_view_list])
		try:
			category_type = st.selectbox(
				"Select Category Type",
				category,
				key=f"category_type {self.key}"	)
			return category_type
		except:
			return StatementConstants.without_category


from django.db import models
from django.contrib.auth.models import PermissionsMixin
from django.contrib.auth.base_user import AbstractBaseUser
from django.utils.translation import ugettext_lazy as _

from users_app.managers import UserManager

class User(AbstractBaseUser, PermissionsMixin):
    username = models.CharField(_('username'), max_length=255, unique=True)
    email = models.EmailField(_('email address'),\
        null=True, blank=True)
    phone = models.CharField(_('phone number'), max_length=30,\
        null=True, blank=True)
    date_joined = models.DateTimeField(_('date joined'), auto_now_add=True)
    is_active = models.BooleanField(_('active'), default=False)
    is_staff = models.BooleanField(_('staff'), default=False)

    is_verified = models.BooleanField(_('verified'), default=False)

    objects = UserManager()

    USERNAME_FIELD = 'username'
    REQUIRED_FIELDS = ['email']

    class Meta:
        verbose_name = _('user')
        verbose_name_plural = _('users')
        unique_together = ('username', 'email', 'phone')


class Solution(object):
    def validMountainArray(self, arr):
      length = len(arr)

      if(length < 3):
         return False

      left = 0
      right = length - 1

      while(True):
        if(left == right or arr[left+1] == arr[left]):
            return False
        if(arr[left+1] > arr[left]):
            left += 1
        elif(arr[left+1] < arr[left]):
           break

      if(left == 0):
         return False

      while(right > left):
         if(arr[right - 1] == arr[right]):
            return False
         if(arr[right - 1] > arr[right]):
            right -= 1
         else:
            return False

      return True

answer = Solution()
print(answer.validMountainArray([9,8,7,6,5,4,3,2,1,0]))


import os
import hydra
import logging

from omegaconf import OmegaConf, DictConfig

# import the dataset loaders
from COIGAN.training.data.datasets_loaders import JsonLineDatasetBase, JsonLineDataset
from COIGAN.training.data.datasets_loaders import SeverstalSteelDefectDataset

# import the dataset generators
from COIGAN.training.data.dataset_generators import SeverstalSteelDefectPreprcessor
from COIGAN.training.data.dataset_generators import ObjectDatasetGenerator
from COIGAN.training.data.dataset_generators import BaseDatasetGenerator
from COIGAN.training.data.dataset_generators import TileDatasetPreprocessor

# import the dataset inspectors
from COIGAN.training.data.dataset_inspectors import JsonLineDatasetInspector
from COIGAN.training.data.dataset_inspectors import JsonLineMaskDatasetInspector

# import the dataset splitters
from COIGAN.training.data.dataset_splitters import BaseSplitter
from COIGAN.training.data.dataset_splitters import FairSplitter

# import the image evaluators
from COIGAN.training.data.image_evaluators import SeverstalBaseEvaluator

# setting up the logger and enable warnings
LOGGER = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)

def check_severstal_jsonl_dataset_health(config: DictConfig):
    """
        Method that check if the jsonl_all dataset
        is correctly created, and if need to be recreated.
    """
    def delete_jsonl_all_dataset():
        LOGGER.info("Severstal jsonl dataset it's damaged, deleting it...")
        os.system(f"rm -r {config.jsonl_dataset_dir}")

    if os.path.exists(config.jsonl_dataset_dir):
        # if the folder already exist do the check
        # if the check is ok return True else
        # delete the folder and return False

        #check if the data folder is present
        if not os.path.exists(
            os.path.join(config.jsonl_dataset_dir, "data")
        ):
            delete_jsonl_all_dataset()
            return False

        # check if the index file is present and not empty
        if not os.path.exists(os.path.join(config.jsonl_dataset_dir, "index")) or \
            not os.path.exists(os.path.join(config.jsonl_dataset_dir, "dataset.jsonl")):
            if os.stat(os.path.join(config.jsonl_dataset_dir, "index")).st_size == 0 or \
                os.stat(os.path.join(config.jsonl_dataset_dir, "dataset.jsonl")).st_size == 0:
                delete_jsonl_all_dataset()
                return False

        # check if the image folder is empty
        if len(os.listdir(os.path.join(config.jsonl_dataset_dir, "data"))) == 0:
            delete_jsonl_all_dataset()
            return False

        #check if the force flag is set
        if config.force_dataset_preprocessor:
            LOGGER.warning("Force dataset preprocessor flag is set, recreating the dataset...")
            delete_jsonl_all_dataset()
            return False

        return True
    else:
        # else return False, because the dataset is not created
        return False

def download_and_extract_severstal_dataset(config: DictConfig):
    """
        Download and extract the Severstal Steel Defect Dataset from kaggle
        only if needed!

    Args:
        dataset_path (str): Path to the dataset
    """

    # check if the extracted dataset is already present if not extract the dataset
    if not os.path.exists(config.raw_dataset_dir):

        # check if the dataset is already downloaded, if not download it
        if not os.path.exists(config.zip_file):
            LOGGER.info(f"Downloading the Steel Defect Dataset in {config.dataset_dir}")
            os.system(f"kaggle competitions download -c severstal-steel-defect-detection -p {config.dataset_dir}")
            # check if the download is completed correctly
            if not os.path.exists(config.zip_file):
                LOGGER.error(f"Dataset can't be downloaded from kaggle!\n \
                    double check the the kaggle api is installed and the kaggle.json file is present in the home directory, \n \
                    if the problem persist downlod the dataset from kaggle and put the zip file in: \n \
                    {config.dataset_dir}")
                raise FileNotFoundError(f"Unable to download the dataset from kaggle in {config.dataset_dir}")
        else:
            LOGGER.info("The dataset is already downloaded")

        # unzip the dataset
        LOGGER.info(f"Unzipping the dataset in {config.raw_dataset_dir}")
        os.system(f"unzip -q {config.zip_file} -d {config.raw_dataset_dir}")
        # delete the zip file
        os.system(f"rm {config.zip_file}")
    else:
        LOGGER.info("The dataset from kaggle is already downloaded and extracted")

def convert_severstal_dataset_to_jsonl(config: DictConfig):
    """
        Convert the Severstal Steel Defect Dataset into COIGAN compatible dataset.

    Args:
        dataset_path (str): Path to the dataset
        output_dir (str): Path to the output directory
    """

    # check if the dataset is already converted and if is converted correctly
    if not check_severstal_jsonl_dataset_health(config):

        # load the dataset and prepare it for the training pipeline
        LOGGER.info("Converting the kaggle dataset in json format...")
        SeverstalSteelDefectPreprcessor(
            SeverstalSteelDefectDataset(config.raw_dataset_dir, mode="all"),
            **config.dataset_preprocessor_kwargs
        ).convert()

    else:
        LOGGER.info("The dataset is already converted in jsonl format!")

def create_severstal_dataset_report(config: DictConfig):
    """
        Create a report of the dataset
    """

    # checking if the report is already generated
    report_dir = os.path.join(config.jsonl_dataset_dir, "reports")
    if os.path.exists(report_dir):
        # check if the force flag is set
        if config.force_dataset_inspector:
            LOGGER.warning("Force dataset inspector flag is set, recreating the dataset report...")

        #check if the raw_report is present and not empty
        elif os.path.exists(os.path.join(report_dir, "raw_report.json")) and \
            os.stat(os.path.join(report_dir, "raw_report.json")).st_size != 0:
                # the report for the next step is preset
                LOGGER.info("The dataset report is already generated!")
                return

        # if the report is present but incomplete or the force flag is set
        # delete it, and create a new one
        os.system(f"rm -r {report_dir}")

    LOGGER.info("Creating the report of the dataset...")

    inspector = JsonLineDatasetInspector(
        **config.dataset_inspector_kwargs
    )

    inspector.inspect()
    inspector.dump_report()
    inspector.dump_raw_report()
    inspector.generate_graphs()

def split_dataset(config: DictConfig):
    """
        Split the dataset in train and test set, keeping the same distribution of the classes
    """

    # checking if the splits are already created or the force flag is set

    train_split_check = os.path.exists(os.path.join(config.train_set_dir))
    test_split_check = os.path.exists(os.path.join(config.test_set_dir))

    if train_split_check and test_split_check:
        # if both are present
        # check if the force flag is set
        if config.force_dataset_splitter:
            LOGGER.warning("Force dataset splitter flag is set, recreating the dataset splits...")
            os.system(f"rm -r {config.dataset_splitter_kwargs.output_dir}")
        else:
            LOGGER.info(f"The train and test splits are already created under the folder {config.dataset_name}! skipping the splitting step...")
            return

    elif train_split_check or test_split_check:
        # if only one is present
        # delete both and recreate them
        LOGGER.info("only one of the splits are present, recreating the splits...")
        os.system(f"rm -r {config.dataset_splitter_kwargs.output_dir}")

    LOGGER.info("Splitting the dataset in train and test set...")

    if config.split_mode == "random":
        LOGGER.info("Splitting the dataset in random mode..")
        BaseSplitter(
            **config.dataset_splitter_kwargs
        ).split()

    elif config.split_mode == "fair":
        LOGGER.info("Splitting the dataset in fair mode..")
        FairSplitter(
            **config.dataset_splitter_kwargs
        ).split()
    else:
        raise ValueError(f"Split mode {config.split_mode} not supported!")

def create_splits_reports(config: DictConfig):

    # checking if the report is already generated
    train_report_dir = os.path.join(config.train_set_dir, "reports")
    test_report_dir = os.path.join(config.test_set_dir, "reports")

    if os.path.exists(train_report_dir) and os.path.exists(test_report_dir):
        # if the report is present skip the process and the force flag is not set
        if config.force_dataset_split_inspector:
            LOGGER.warning("Force dataset split inspector flag is set, deletting the train and test splits report...")
            os.system(f"rm -r {train_report_dir}")
            os.system(f"rm -r {test_report_dir}")

        else:
            LOGGER.info("The train split report is already generated! skipping the report generation step...")
            return

    LOGGER.info("Creating the report for the train and test splits...")

    inspector = JsonLineDatasetInspector(
        **config.train_split_inspector_kwargs
    )

    inspector.inspect()
    inspector.dump_report()
    inspector.dump_raw_report()
    inspector.generate_graphs()

    inspector = JsonLineDatasetInspector(
        **config.test_split_inspector_kwargs
    )

    inspector.inspect()
    inspector.dump_report()
    inspector.dump_raw_report()
    inspector.generate_graphs()

def create_object_datasets(config: DictConfig):

    # checking if the object datasets are already created
    if config.force_object_dataset_generator:
        LOGGER.warning("Force object dataset generator flag is set, deletting the object datasets...")
        if os.path.exists(config.object_datasets_dir):
            os.system(f"rm -r {config.object_datasets_dir}")

    else:
        count = 0
        for _class in config.object_target_classes:
            if os.path.exists(os.path.join(config.object_datasets_dir, f"object_dataset_{_class}")):
                count += 1
        if count == len(config.object_target_classes):
            LOGGER.info("All the object datasets are already created!")
            return
        elif count > 0:
            LOGGER.info("Some object datasets are already created! deleting them...")
            os.system(f"rm -r {config.object_datasets_dir}")
        # if count == 0 the object datasets are not created, so we can continue

    # loading the input dataset
    dataset = JsonLineDatasetBase(
        os.path.join(config.train_set_dir, "dataset.jsonl"),
        os.path.join(config.train_set_dir, "index"),
        config.binary
    )
    dataset.on_worker_init()

    # creating one object dataset for each class
    for _class in config.object_target_classes:
        LOGGER.info(f"Creating the object dataset for the class {_class}...")
        ObjectDatasetGenerator(
            input_dataset=dataset,
            image_dir=os.path.join(config.train_set_dir, "data"),
            target_class=_class,
            **config.object_dataset_generator_kwargs,
        ).convert()

def create_object_datasets_reports(config: DictConfig):
    """
        Create the reports for the object datasets
    """
    if config.force_object_dataset_inspector:
        LOGGER.warning("Force object dataset inspector flag is set, deletting the object datasets reports...")
        for _class in config.object_target_classes:
            report_dir = os.path.join(config.object_datasets_dir, f"object_dataset_{_class}", "reports")
            if os.path.exists(report_dir):
                os.system(f"rm -r {report_dir}")

    else:
        # checking if the report is already generated
        reports_found = 0
        for _class in config.object_target_classes:
            report_dir = os.path.join(config.object_datasets_dir, f"object_dataset_{_class}", "reports")
            if os.path.exists(report_dir) and \
                os.path.exists(os.path.join(report_dir, "brief_report.json")):
                    reports_found += 1

        if reports_found == len(config.object_target_classes):
            # if the report is present skip the process
            LOGGER.info("The object datasets reports are already generated! skipping the report generation step...")
            return

        elif reports_found > 0:
            LOGGER.info("Some object datasets reports are already generated! deleting them...")
            for _class in config.object_target_classes:
                report_dir = os.path.join(config.object_datasets_dir, f"object_dataset_{_class}", "reports")
                if os.path.exists(report_dir):
                    os.system(f"rm -r {report_dir}")

        # if reports_found == 0 the object datasets reports are not created, so we can continue

    LOGGER.info("Creating the reports for the object datasets...")
    for _class in config.object_target_classes:
        LOGGER.info(f"Creating the report for the object dataset {_class}...")
        inspector = JsonLineMaskDatasetInspector(
            os.path.join(config.object_datasets_dir, f"object_dataset_{_class}"),
            os.path.join(config.object_datasets_dir, f"object_dataset_{_class}", "reports"),
            config.binary
        )

        inspector.inspect()
        inspector.dump_report()
        inspector.dump_raw_report()
        inspector.generate_graphs()

def create_base_dataset(config: DictConfig):
    """
        Create the base dataset
    """
    if config.force_base_dataset_generator:
        LOGGER.warning("Force base dataset generator flag is set, deletting the base dataset...")
        if os.path.exists(config.base_dataset_dir):
            os.system(f"rm -r {config.base_dataset_dir}")

    elif os.path.exists(config.base_dataset_dir):
        # checking if the base dataset is empty
        if len(os.listdir(os.path.join(config.base_dataset_dir, "data"))) == 0:
            LOGGER.info("The base dataset is empty! deleting it...")
            os.system(f"rm -r {config.base_dataset_dir}")
        else:
            LOGGER.info("The base dataset is already created! skipping the base dataset creation step...")
            return

    # load the input dataset
    input_dataset = JsonLineDataset(
        image_folder_path =     os.path.join(config.train_set_dir, "data"),
        metadata_file_path =    os.path.join(config.train_set_dir, "dataset.jsonl"),
        index_file_path =       os.path.join(config.train_set_dir, "index"),
        masks_fields =          ["polygons"],
        classes =               config.object_target_classes,
        size =                  config.original_tile_size,
        binary =                config.binary
    )

    # initialize the input dataset
    input_dataset.on_worker_init()

    # load the image evaluator
    img_evaluator = SeverstalBaseEvaluator(
        **config.base_evaluator_kwargs
    )

    LOGGER.info("Creating the base dataset...")
    BaseDatasetGenerator(
        input_dataset =     input_dataset,
        image_dir=          os.path.join(config.train_set_dir, "data"),
        img_evaluator=      img_evaluator,
        **config.base_dataset_generator_kwargs
    ).convert()

def create_tile_datasets(config: DictConfig):
    """
    Method that create the tile datasets

    Args:
        config (DictConfig): _description_
    """
    def delete_tile_datasets(config):
        if os.path.exists(config.tile_train_set_dir):
            LOGGER.info("Deleting the tile train set...")
            os.system(f"rm -r {config.tile_train_set_dir}")
        if os.path.exists(config.tile_test_set_dir):
            LOGGER.info("Deleting the tile test set...")
            os.system(f"rm -r {config.tile_test_set_dir}")

    if config.force_tile_dataset_generator:
        LOGGER.warning("Force tile dataset generator flag is set, deletting the tile datasets...")
        delete_tile_datasets(config)
    elif (not os.path.exists(config.tile_train_set_dir)) or (not os.path.exists(config.tile_test_set_dir)):
        LOGGER.info("The tile datasets are not created! creating them...")
        delete_tile_datasets(config)
    elif os.path.exists(config.tile_train_set_dir) and os.path.exists(config.tile_test_set_dir):
        LOGGER.info("The tile datasets are already created! skipping the tile dataset creation step...")
        return

    # load the train dataset
    train_dataset = JsonLineDataset(
        image_folder_path =     os.path.join(config.train_set_dir, "data"),
        metadata_file_path =    os.path.join(config.train_set_dir, "dataset.jsonl"),
        index_file_path =       os.path.join(config.train_set_dir, "index"),
        masks_fields =          ["polygons"],
        classes =               config.object_target_classes,
        size =                  config.original_tile_size,
        binary =                config.binary
    )

    # load the test dataset
    test_dataset = JsonLineDataset(
        image_folder_path =     os.path.join(config.test_set_dir, "data"),
        metadata_file_path =    os.path.join(config.test_set_dir, "dataset.jsonl"),
        index_file_path =       os.path.join(config.test_set_dir, "index"),
        masks_fields =          ["polygons"],
        classes =               config.object_target_classes,
        size =                  config.original_tile_size,
        binary =                config.binary
    )

    # initialize the datasets
    train_dataset.on_worker_init()
    test_dataset.on_worker_init()

        # load the image evaluator
    image_evaluator = SeverstalBaseEvaluator(
        **config.base_evaluator_kwargs
    )

    LOGGER.info("Creating the tile train dataset...")
    TileDatasetPreprocessor(
        input_dataset = train_dataset,
        image_evaluator = image_evaluator,
        **config.tile_train_dataset_generator_kwargs
    ).convert()

    LOGGER.info("Creating the tile test dataset...")
    TileDatasetPreprocessor(
        input_dataset = test_dataset,
        **config.tile_test_dataset_generator_kwargs
    ).convert()

@hydra.main(config_path='../configs/data_preparation', config_name='severstal_dataset_preparation.yaml', version_base="1.1")
def main(config: DictConfig):
    """
        Script that load the Severstal steel defect dataset and prepare it for the trianing pipeline.

        Operations:
            V 1 - Download the dataset from kaggle
            V 2 - Extract the dataset and delete the zip file
            V 3 - Convert the dataset in jsonl format
                V 3.1 - create a report of the dataset
            V 4 - Split the dataset in train, and test (keeping the same distribution of the classes)
                V 4.1 - create a report of the train and test splits
            V 5 - Create the defects datasets to train the COIGAN model to inpaint the defects and stylegan to generate new masks, from the train dataset
                V 5.1 - Create the reports for the object datasets
            V 6 - Create the none defected img dataset to train another stylegan, to generate new none defected images, or to use directly in the inpainting process.
            7 - Convert in tile datasets the train and the test dataset, with the same size of the tiles in the precedenet processes

        TODO: list of bugs that need to be fixed or improvements:

    """

    LOGGER.info(f"Current folder: {os.getcwd()}")

    # 1 & 2 - download and extract the dataset if necessary
    download_and_extract_severstal_dataset(config)

    # 3 - convert the dataset in jsonl format
    convert_severstal_dataset_to_jsonl(config)

    # 3.1 - create report of the dataset
    create_severstal_dataset_report(config)

    # 4 - split the dataset in train and test set
    split_dataset(config)

    # 4.1 - create reports for the splitted datasets
    create_splits_reports(config)

    # 5 - Create the defects datasets to train the COIGAN model to inpaint the defects and stylegan to generate new masks, from the train dataset
    create_object_datasets(config)

    # 5.1 - create the reports for the object datasets
    create_object_datasets_reports(config)

    # 6 - create the none defected img dataset to train another stylegan, to generate new none defected images
    create_base_dataset(config)

    # 7 - convert in tile datasets the train and the test dataset, with the same size of the tiles used in the COIGAN model
    create_tile_datasets(config)

if __name__ == "__main__":
    main()


import os

from qiskit import QuantumCircuit
from environments.circuits import QubitCircuit

def import_test_set():
    print('Importing test set...')

    directory_path = "./realistic_test_set/"

    files = os.listdir(directory_path)
    qasm_files = list(filter(lambda file_name: len(file_name) > 5 and file_name[-5:] == ".qasm", files))

    circuits = []

    for i,file_name in enumerate(qasm_files):
        file_path = directory_path + file_name

        if os.path.getsize(file_path) > 10000:
            continue

        qiskit_circuit = QuantumCircuit.from_qasm_file(file_path)

        gates = []

        for gate_obj, qubits, _ in qiskit_circuit.data:
            if len(qubits) > 1:
                if gate_obj.__class__.__name__ not in  ["CnotGate", "CXGate"]:
                    exit("Non-cnot gate (" + gate_obj.__class__.__name__ + ") found for circuit: " + str(file_name))

                gate = (qubits[0].index, qubits[1].index)
                gates.append(gate)

        circuit = QubitCircuit.from_gates(16, gates)

        circuits.append(circuit)

        if i % int(len(qasm_files) / 10) == 0:
            print('Import ' + str(int(100 * i / len(qasm_files))) + '% complete')

    return list(filter(lambda c: c.depth() < 200, circuits))

# circuits = list(filter(lambda c: c.depth() < 100, import_test_set()))

# for circuit in circuits:
#     print('Circuit depth:', circuit.depth())
#     max_qubits = [max(q1,q2) for (q1,q2) in circuit.gates]
#     print('Max qubit used:', max(max_qubits))
#     print()


from ._input import Input

class CheckBoxInput(Input):
    def __init__(
        self, title: str, checkboxes: list, opt_value: str, optional: bool = True
    ) -> None:
        super(CheckBoxInput, self).__init__()
        _type = "checkboxes"
        self.checkboxes = checkboxes
        _options = self.get_opts(opt_value)
        _element = dict(type=_type, options=_options, action_id="checkboxes-action")
        _label = dict(type="plain_text", text=title, emoji=True)

        if not optional:
            self.body = dict(type=self.type, element=_element, label=_label)
        else:
            self.body = dict(
                type=self.type, element=_element, label=_label, optional=optional
            )

    def get_opts(self, value: str) -> list:
        option_list = [
            dict(
                text=dict(type="plain_text", text=checkbox, emoji=True),
                value=f"{value}-{idx}",
            )
            for idx, checkbox in enumerate(self.checkboxes)
        ]
        return option_list


import unittest
from itertools import chain
from typing import Dict, Optional, Tuple
from unittest.mock import MagicMock, patch

import pytest
from confluent_kafka.schema_registry.schema_registry_client import (
    RegisteredSchema,
    Schema,
)

from datahub.configuration.common import ConfigurationError
from datahub.emitter.mce_builder import (
    make_dataplatform_instance_urn,
    make_dataset_urn_with_platform_instance,
)
from datahub.ingestion.api.common import PipelineContext
from datahub.ingestion.api.workunit import MetadataWorkUnit
from datahub.ingestion.source.kafka import KafkaSource, KafkaSourceConfig
from datahub.metadata.com.linkedin.pegasus2avro.mxe import MetadataChangeEvent
from datahub.metadata.schema_classes import (
    BrowsePathsClass,
    DataPlatformInstanceClass,
    KafkaSchemaClass,
    SchemaMetadataClass,
)

class KafkaSourceTest(unittest.TestCase):
    @patch("datahub.ingestion.source.kafka.confluent_kafka.Consumer", autospec=True)
    def test_kafka_source_configuration(self, mock_kafka):
        ctx = PipelineContext(run_id="test")
        kafka_source = KafkaSource(
            KafkaSourceConfig.parse_obj({"connection": {"bootstrap": "foobar:9092"}}),
            ctx,
        )
        kafka_source.close()
        assert mock_kafka.call_count == 1

    @patch("datahub.ingestion.source.kafka.confluent_kafka.Consumer", autospec=True)
    def test_kafka_source_workunits_wildcard_topic(self, mock_kafka):
        mock_kafka_instance = mock_kafka.return_value
        mock_cluster_metadata = MagicMock()
        mock_cluster_metadata.topics = ["foobar", "bazbaz"]
        mock_kafka_instance.list_topics.return_value = mock_cluster_metadata

        ctx = PipelineContext(run_id="test")
        kafka_source = KafkaSource(
            KafkaSourceConfig.parse_obj(
                {"connection": {"bootstrap": "localhost:9092"}}
            ),
            ctx,
        )
        workunits = list(kafka_source.get_workunits())

        first_mce = workunits[0].metadata
        assert isinstance(first_mce, MetadataChangeEvent)
        mock_kafka.assert_called_once()
        mock_kafka_instance.list_topics.assert_called_once()
        assert len(workunits) == 4

    @patch("datahub.ingestion.source.kafka.confluent_kafka.Consumer", autospec=True)
    def test_kafka_source_workunits_topic_pattern(self, mock_kafka):
        mock_kafka_instance = mock_kafka.return_value
        mock_cluster_metadata = MagicMock()
        mock_cluster_metadata.topics = ["test", "foobar", "bazbaz"]
        mock_kafka_instance.list_topics.return_value = mock_cluster_metadata

        ctx = PipelineContext(run_id="test1")
        kafka_source = KafkaSource.create(
            {
                "topic_patterns": {"allow": ["test"]},
                "connection": {"bootstrap": "localhost:9092"},
            },
            ctx,
        )
        workunits = [w for w in kafka_source.get_workunits()]

        mock_kafka.assert_called_once()
        mock_kafka_instance.list_topics.assert_called_once()
        assert len(workunits) == 2

        mock_cluster_metadata.topics = ["test", "test2", "bazbaz"]
        ctx = PipelineContext(run_id="test2")
        kafka_source = KafkaSource.create(
            {
                "topic_patterns": {"allow": ["test.*"]},
                "connection": {"bootstrap": "localhost:9092"},
            },
            ctx,
        )
        workunits = [w for w in kafka_source.get_workunits()]
        assert len(workunits) == 4

    @patch("datahub.ingestion.source.kafka.confluent_kafka.Consumer", autospec=True)
    def test_kafka_source_workunits_with_platform_instance(self, mock_kafka):

        PLATFORM_INSTANCE = "kafka_cluster"
        PLATFORM = "kafka"
        TOPIC_NAME = "test"

        mock_kafka_instance = mock_kafka.return_value
        mock_cluster_metadata = MagicMock()
        mock_cluster_metadata.topics = [TOPIC_NAME]
        mock_kafka_instance.list_topics.return_value = mock_cluster_metadata

        ctx = PipelineContext(run_id="test1")
        kafka_source = KafkaSource.create(
            {
                "connection": {"bootstrap": "localhost:9092"},
                "platform_instance": PLATFORM_INSTANCE,
            },
            ctx,
        )
        workunits = [w for w in kafka_source.get_workunits()]

        # We should only have 1 topic + sub-type wu.
        assert len(workunits) == 2
        assert isinstance(workunits[0], MetadataWorkUnit)
        assert isinstance(workunits[0].metadata, MetadataChangeEvent)
        proposed_snap = workunits[0].metadata.proposedSnapshot
        assert proposed_snap.urn == make_dataset_urn_with_platform_instance(
            platform=PLATFORM,
            name=TOPIC_NAME,
            platform_instance=PLATFORM_INSTANCE,
            env="PROD",
        )

        # DataPlatform aspect should be present when platform_instance is configured
        data_platform_aspects = [
            asp
            for asp in proposed_snap.aspects
            if type(asp) == DataPlatformInstanceClass
        ]
        assert len(data_platform_aspects) == 1
        assert data_platform_aspects[0].instance == make_dataplatform_instance_urn(
            PLATFORM, PLATFORM_INSTANCE
        )

        # The default browse path should include the platform_instance value
        browse_path_aspects = [
            asp for asp in proposed_snap.aspects if type(asp) == BrowsePathsClass
        ]
        assert len(browse_path_aspects) == 1
        assert (
            f"/prod/{PLATFORM}/{PLATFORM_INSTANCE}/{TOPIC_NAME}"
            in browse_path_aspects[0].paths
        )

    @patch("datahub.ingestion.source.kafka.confluent_kafka.Consumer", autospec=True)
    def test_close(self, mock_kafka):
        mock_kafka_instance = mock_kafka.return_value
        ctx = PipelineContext(run_id="test")
        kafka_source = KafkaSource.create(
            {
                "topic_patterns": {"allow": ["test.*"]},
                "connection": {"bootstrap": "localhost:9092"},
            },
            ctx,
        )
        kafka_source.close()
        assert mock_kafka_instance.close.call_count == 1

    def test_kafka_source_stateful_ingestion_requires_platform_instance(
        self,
    ):
        class StatefulProviderMock:
            def __init__(self, config, ctx):
                self.ctx = ctx
                self.config = config

            def is_stateful_ingestion_configured(self):
                return self.config.stateful_ingestion.enabled

        kafka_source_patcher = unittest.mock.patch.object(
            KafkaSource, "__bases__", (StatefulProviderMock,)
        )

        ctx = PipelineContext(run_id="test", pipeline_name="test")
        with pytest.raises(ConfigurationError):
            with kafka_source_patcher:
                # prevent delattr on __bases__ on context __exit__
                kafka_source_patcher.is_local = True
                KafkaSource.create(
                    {
                        "stateful_ingestion": {"enabled": "true"},
                        "connection": {"bootstrap": "localhost:9092"},
                    },
                    ctx,
                )

    @patch(
        "datahub.ingestion.source.kafka.confluent_kafka.schema_registry.schema_registry_client.SchemaRegistryClient",
        autospec=True,
    )
    @patch("datahub.ingestion.source.kafka.confluent_kafka.Consumer", autospec=True)
    def test_kafka_source_workunits_schema_registry_subject_name_strategies(
        self, mock_kafka_consumer, mock_schema_registry_client
    ):
        # Setup the topic to key/value schema mappings for all types of schema registry subject name strategies.
        # <key=topic_name, value=(<key_schema>,<value_schema>)
        topic_subject_schema_map: Dict[
            str, Tuple[RegisteredSchema, RegisteredSchema]
        ] = {
            # TopicNameStrategy is used for subject
            "topic1": (
                RegisteredSchema(
                    schema_id="schema_id_2",
                    schema=Schema(
                        schema_str='{"type":"record", "name":"Topic1Key", "namespace": "test.acryl", "fields": [{"name":"t1key", "type": "string"}]}',
                        schema_type="AVRO",
                    ),
                    subject="topic1-key",
                    version=1,
                ),
                RegisteredSchema(
                    schema_id="schema_id_1",
                    schema=Schema(
                        schema_str='{"type":"record", "name":"Topic1Value", "namespace": "test.acryl", "fields": [{"name":"t1value", "type": "string"}]}',
                        schema_type="AVRO",
                    ),
                    subject="topic1-value",
                    version=1,
                ),
            ),
            # RecordNameStrategy is used for subject
            "topic2": (
                RegisteredSchema(
                    schema_id="schema_id_3",
                    schema=Schema(
                        schema_str='{"type":"record", "name":"Topic2Key", "namespace": "test.acryl", "fields": [{"name":"t2key", "type": "string"}]}',
                        schema_type="AVRO",
                    ),
                    subject="test.acryl.Topic2Key",
                    version=1,
                ),
                RegisteredSchema(
                    schema_id="schema_id_4",
                    schema=Schema(
                        schema_str='{"type":"record", "name":"Topic2Value", "namespace": "test.acryl", "fields": [{"name":"t2value", "type": "string"}]}',
                        schema_type="AVRO",
                    ),
                    subject="test.acryl.Topic2Value",
                    version=1,
                ),
            ),
            # TopicRecordNameStrategy is used for subject
            "topic3": (
                RegisteredSchema(
                    schema_id="schema_id_4",
                    schema=Schema(
                        schema_str='{"type":"record", "name":"Topic3Key", "namespace": "test.acryl", "fields": [{"name":"t3key", "type": "string"}]}',
                        schema_type="AVRO",
                    ),
                    subject="topic3-test.acryl.Topic3Key-key",
                    version=1,
                ),
                RegisteredSchema(
                    schema_id="schema_id_5",
                    schema=Schema(
                        schema_str='{"type":"record", "name":"Topic3Value", "namespace": "test.acryl", "fields": [{"name":"t3value", "type": "string"}]}',
                        schema_type="AVRO",
                    ),
                    subject="topic3-test.acryl.Topic3Value-value",
                    version=1,
                ),
            ),
        }

        # Mock the kafka consumer
        mock_kafka_instance = mock_kafka_consumer.return_value
        mock_cluster_metadata = MagicMock()
        mock_cluster_metadata.topics = list(topic_subject_schema_map.keys())
        mock_cluster_metadata.topics.append("schema_less_topic")
        mock_kafka_instance.list_topics.return_value = mock_cluster_metadata

        # Mock the schema registry client
        # - mock get_subjects: all subjects in topic_subject_schema_map
        mock_schema_registry_client.return_value.get_subjects.return_value = [
            v.subject for v in chain(*topic_subject_schema_map.values())
        ]

        # - mock get_latest_version
        def mock_get_latest_version(subject_name: str) -> Optional[RegisteredSchema]:
            for registered_schema in chain(*topic_subject_schema_map.values()):
                if registered_schema.subject == subject_name:
                    return registered_schema
            return None

        mock_schema_registry_client.return_value.get_latest_version = (
            mock_get_latest_version
        )

        # Test the kafka source
        source_config = {
            "connection": {"bootstrap": "localhost:9092"},
            # Setup the topic_subject_map for topic2 which uses RecordNameStrategy
            "topic_subject_map": {
                "topic2-key": "test.acryl.Topic2Key",
                "topic2-value": "test.acryl.Topic2Value",
            },
        }
        ctx = PipelineContext(run_id="test")
        kafka_source = KafkaSource.create(source_config, ctx)
        workunits = list(kafka_source.get_workunits())

        mock_kafka_consumer.assert_called_once()
        mock_kafka_instance.list_topics.assert_called_once()
        assert len(workunits) == 8
        i: int = -1
        for wu in workunits:
            assert isinstance(wu, MetadataWorkUnit)
            if not isinstance(wu.metadata, MetadataChangeEvent):
                continue
            mce: MetadataChangeEvent = wu.metadata
            i += 1

            if i < len(topic_subject_schema_map.keys()):
                # First 3 workunits (topics) must have schemaMetadata aspect
                assert isinstance(mce.proposedSnapshot.aspects[1], SchemaMetadataClass)
                schemaMetadataAspect: SchemaMetadataClass = (
                    mce.proposedSnapshot.aspects[1]
                )
                assert isinstance(schemaMetadataAspect.platformSchema, KafkaSchemaClass)
                # Make sure the schema name is present in topic_subject_schema_map.
                assert schemaMetadataAspect.schemaName in topic_subject_schema_map
                # Make sure the schema_str matches for the key schema.
                assert (
                    schemaMetadataAspect.platformSchema.keySchema
                    == topic_subject_schema_map[schemaMetadataAspect.schemaName][
                        0
                    ].schema.schema_str
                )
                # Make sure the schema_str matches for the value schema.
                assert (
                    schemaMetadataAspect.platformSchema.documentSchema
                    == topic_subject_schema_map[schemaMetadataAspect.schemaName][
                        1
                    ].schema.schema_str
                )
                # Make sure we have 2 fields, one from the key schema & one from the value schema.
                assert len(schemaMetadataAspect.fields) == 2
            else:
                # Last topic('schema_less_topic') has no schema defined in the registry.
                # The schemaMetadata aspect should not be present for this.
                for aspect in mce.proposedSnapshot.aspects:
                    assert not isinstance(aspect, SchemaMetadataClass)


import os
import json
import shutil
import numpy as np
from tqdm import tqdm
import copy

#########################################################
# 多个文件：yolo格式 转 labelme格式
#########################################################
yolo_dict_format = {   # 初始化labelme格式
    "version": "5.1.1",                 # labelme版本
    "flags": {},
    "shapes": [{                 # 标注
        "label": "people",              # 类别：如关键点类别left_eye
        "points": [[0, 0], [0, 0]],     # 关键点坐标
        "group_id": 0,                  # 组的id
        "shape_type": "rectangle",      # 标注类别：如point，rectangle
        "flags": {}
    }],
    "imagePath": "00004.png",
    "imageData": None,
    "imageHeight": 1080,
    "imageWidth": 1920
}
keypoints = ['nose', 'left_eye', 'right_eye', 'left_ear', 'right_ear',
                  'left_shoulder', 'right_shoulder', 'left_elbow', 'right_elbow', 'left_wrist',
                  'right_wrist', 'left_hip', 'right_hip', 'left_knee', 'right_knee',
                  'left_ankle', 'right_ankle']

yolo_txt_path = 'E:\\dataset\\set01\\runs\pose\\predict2\\labels\\'    # yolo格式txt文件路径
labelme_json_path = 'E:\\dataset\\set01\\runs\pose\\predict2\\json\\'  # labelme格式json文件路径

extension = '.txt'      # 文件后缀
files = [file for file in os.listdir(yolo_txt_path) if file.endswith(extension)]    # 查看指定目录下指定后缀的列表
for file in tqdm(files):
    print(file)
    yolo_txt = yolo_txt_path + file
    with open(yolo_txt, 'r', encoding='utf-8') as f:
        lines = f.readlines()       # 按行读取
        print(lines)
        yolo_dict = copy.deepcopy(yolo_dict_format)    # 一张图片一个字典
        yolo_dict["imagePath"] = str(int(file.split('.')[0].split('_')[-1]) - 0).zfill(5) + '.png'
        people_id = 0
        for line in lines:          # 遍历每一行
            line = line.replace("\n", "")
            line_list = line.split(' ')
            print(line_list)
            for i in range(len(line_list)):
                if i == 0:
                    yolo_dict["shapes"][-1]["label"] = "people"
                    yolo_dict["shapes"][-1]["shape_type"] = "rectangle"
                    yolo_dict["shapes"][-1]["group_id"] = people_id
                elif i == 1:
                    center_x = float(line_list[i]) * yolo_dict["imageWidth"]
                elif i == 2:
                    center_y = float(line_list[i]) * yolo_dict["imageHeight"]
                elif i == 3:
                    yolo_dict["shapes"][-1]["points"] = [[0, 0], [0, 0]]
                    bbox_width = float(line_list[i]) * yolo_dict["imageWidth"]
                    yolo_dict["shapes"][-1]["points"][0][0] = center_x - bbox_width / 2
                    yolo_dict["shapes"][-1]["points"][1][0] = center_x + bbox_width / 2
                elif i == 4:
                    bbox_height = float(line_list[i]) * yolo_dict["imageHeight"]
                    yolo_dict["shapes"][-1]["points"][0][1] = center_y - bbox_height / 2
                    yolo_dict["shapes"][-1]["points"][1][1] = center_y + bbox_height / 2
                elif i % 2 == 1:
                    yolo_dict["shapes"].append(dict())
                    yolo_dict["shapes"][-1]["group_id"] = people_id
                    yolo_dict["shapes"][-1]["flags"] = {}
                    yolo_dict["shapes"][-1]["label"] = keypoints[int(i/2-2.5)]
                    yolo_dict["shapes"][-1]["shape_type"] = "point"
                    yolo_dict["shapes"][-1]["points"] = [[0, 0]]
                    yolo_dict["shapes"][-1]["points"][0][0] = float(line_list[i]) * yolo_dict["imageWidth"]
                elif i % 2 == 0:
                    yolo_dict["shapes"][-1]["points"][0][1] = float(line_list[i]) * yolo_dict["imageHeight"]
            people_id += 1
    print(yolo_dict)
    labelme_json = labelme_json_path + yolo_dict["imagePath"].split('.')[-2] + '.json'
    with open(labelme_json, "w", encoding='utf-8') as f:
        json.dump(yolo_dict, f)