"""
Custom converters for the flask router.
"""
import re
from werkzeug.routing import BaseConverter

class ListConverter(BaseConverter):

    def to_python(self, value):
        return value.split(',')

    def to_url(self, values):
        return ','.join(BaseConverter.to_url(value) for value in values)

class LuceneQueryConverter(BaseConverter):
    """
    Quotes and unquote Lucene search strings.
    """

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)

        lucene_special_chars = r'+-&|!(){}[]^"~*?:\\'
        self.pattern = re.compile(
            f"""# verbose regex, ignores whitespace during compilation
            (   # Capturing group #1

                # Match any of the lucene special characters exactly once. The chars
                # contain many regex special chars which need to be escaped.
                [{re.escape(lucene_special_chars)}] # Whitespace isn't ignored in [character clauses]
            )
            """,
            re.VERBOSE
        )

    def to_python(self, text: str) -> str:
        """
        Escapes lucene special characters.

        Here's what Lucene says needs to be escaped: `+ - && || ! ( ) { } [ ] ^ " ~ * ? : \`
        From https://lucene.apache.org/core/2_9_4/queryparsersyntax.html#Escaping%20Special%20Characters

        It's not clear whether && (and ||) should be escaped as \&& or \&\&, we're going with the latter.
        """

        replacement = r'\\\1'
        (quoted, number_of_subs_made) = re.subn(self.pattern, replacement, text)
        if self.pattern.match(text):
            print(self.pattern.match(text).groups())
        print(text)
        print(quoted)
        print(number_of_subs_made)
        return quoted

    def to_url(self, text: str):
        """
        Remove quotes added to escape lucene special characters.
        """
        unquoted = re.sub(r'"\\([\+\-!\(\)\{\}\[\]\^\"~\*\?\:\\/&\|])"', r'\1', text)
        return super().to_url(unquoted)

class ReviewServiceConverter(BaseConverter):
    """
    Maps review service names provided as parameter in the url with their internal names used in the database.
    """

    review_services = {
        'reviewcommons': 'review commons',
        'review commons': 'review commons',
        'elife': 'elife',
        'embopress': 'embo press',
        'embo press': 'embo press',
        'rrc19': 'MIT Press - Journals',
        'mit press - journals': 'MIT Press - Journals',
        'pci': 'Peer Community In',
        'peer community in': 'Peer Community In',
        'peerageofscience': 'peerage of science',
        'peerage of science': 'peerage of science',
    }

    def to_python(self, text: str) -> str:
        """
        Returns the internal review service name.
        """
        review_service_name = self.review_services[text.lower()]
        return review_service_name


from pyecharts.charts import Bar
from pyecharts import options as opts
import pandas as pd

main_file_location = 'Target1.csv'

#--------------------access dataset and clear unuse data-----------------------------#
data1 = pd.read_csv(main_file_location)
# replace NaN value with 'Unspecified' in the last six columns and give it to 'data1_clearNaN'
data1_clearNaN = data1[['Critic_Score','Critic_Count','User_Score','User_Count','Developer','Rating']].fillna('Unspecified')
# Use data1_clearNaN to mearge the original data's last six columns,so the new data will only have NaN value in 'Year_of_Release' column:
data1[['Critic_Score','Critic_Count','User_Score','User_Count','Developer','Rating']] = data1_clearNaN
# drop rows that don't have a specific year.
data1 = data1.dropna(axis=0)
# after dropping NaN values, I can change this column's type to int:
data1['Year_of_Release'] = data1['Year_of_Release'].astype('int')
# drop data belongs to year 2020 and year 2017
# because there are only 4 games in these two year in this dataset,and also the aim is before year 2016.
list1 = list(data1['Year_of_Release'].unique())
list1.remove(2020)
list1.remove(2017)
data1 = data1[data1['Year_of_Release'].isin(list1)]
#------------------------------end clearing-------------------------------------#

v1 = data1[['Year_of_Release','NA_Sales','EU_Sales','JP_Sales','Other_Sales','Global_Sales',]].groupby('Year_of_Release').sum().sort_values(by='Year_of_Release',ascending=False)
v1 = v1.reset_index().rename(columns = {'Year_of_Release':'Year'})

Year_list=list(v1['Year'])
NA_list=list(v1['NA_Sales'])
EU_list=list(v1['EU_Sales'])
JP_list=list(v1['JP_Sales'])
Other_list=list(v1['Other_Sales'])
Global_list=list(v1['Global_Sales'])

def two_decimal(x):
    for i in range(len(x)):
        a = x[i]
        a = round(a,2)
        x[i] = a
    return x

two_decimal(NA_list)
two_decimal(EU_list)
two_decimal(JP_list)
two_decimal(Other_list)
two_decimal(Global_list)

bar1 = (
       Bar()
       .add_xaxis(Year_list)
       .add_yaxis('NA Sales',NA_list,stack='stack1')
       .add_yaxis('EU Sales',EU_list,stack='stack1')
       .add_yaxis('JP Sales',JP_list,stack='stack1')
       .add_yaxis('Other Sales',Other_list,stack='stack1')
       .set_global_opts(
               title_opts=opts.TitleOpts(
                       title="Each Year's Total Sales in Areas(millions of units)",
                       pos_left='center'
                       ),
               datazoom_opts=opts.DataZoomOpts(),
               legend_opts=opts.LegendOpts(
                       pos_right = 100,
                       pos_top = 70,
                       orient = 'vertical'
                       ),
               )
       .set_series_opts(
               label_opts=opts.LabelOpts(
                       is_show=False,
                       )
               )
       .render("Each Year's Total Sales in Areas.html")
       )

bar2 = (
        Bar()
        .add_xaxis(Year_list)
        .add_yaxis('Global Sales',Global_list,category_gap='35%')
        .set_global_opts(
               title_opts=opts.TitleOpts(
                       title="Total global sales per year(millions of units)",
                       pos_left='center'
                       ),
               legend_opts=opts.LegendOpts(
                       pos_right = 100,
                       pos_top = 70,
                       orient = 'vertical'
                       ),
               datazoom_opts=opts.DataZoomOpts(),
               )
        .render("Total global sales per year.html")
        )


from uuid import uuid4

from django.core.exceptions import ValidationError
from django.core.validators import int_list_validator
from django.db import models
from django.db.models import Q
from django.utils.translation import gettext_lazy as _

from django_ledger.models import CreateUpdateMixIn
from django_ledger.models.utils import LazyLoader

lazy_loader = LazyLoader()

class BankAccountModelManager(models.Manager):

    def for_entity(self, entity_slug: str, user_model):
        qs = self.get_queryset()
        return qs.filter(
            Q(ledger__entity__slug__exact=entity_slug) &
            (
                    Q(ledger__entity__admin=user_model) |
                    Q(ledger__entity__managers__in=[user_model])
            )
        )

class BackAccountModelAbstract(CreateUpdateMixIn):
    REL_NAME_PREFIX = 'bank'
    ACCOUNT_TYPES = [
        ('checking', _('Checking')),
        ('savings', _('Savings')),
        ('money_mkt', _('Money Market')),
    ]

    uuid = models.UUIDField(default=uuid4, editable=False, primary_key=True)
    name = models.CharField(max_length=150, null=True, blank=True)
    account_number = models.CharField(max_length=30, null=True, blank=True,
                                      validators=[
                                          int_list_validator(sep='', message=_('Only digits allowed'))
                                      ])
    routing_number = models.CharField(max_length=30, null=True, blank=True,
                                      validators=[
                                          int_list_validator(sep='', message=_('Only digits allowed'))
                                      ])
    aba_number = models.CharField(max_length=30, null=True, blank=True)
    account_type = models.CharField(choices=ACCOUNT_TYPES, max_length=10)
    cash_account = models.ForeignKey('django_ledger.AccountModel',
                                     on_delete=models.CASCADE,
                                     verbose_name=_('Cash Account'),
                                     related_name=f'{REL_NAME_PREFIX}_cash_account',
                                     null=True, blank=True)
    ledger = models.OneToOneField('django_ledger.LedgerModel',
                                  editable=False,
                                  verbose_name=_('Ledger'),
                                  on_delete=models.CASCADE)
    active = models.BooleanField(default=True)
    hidden = models.BooleanField(default=False)
    objects = BankAccountModelManager()

    def configure(self,
                  entity_slug,
                  user_model,
                  posted_ledger: bool = True):
        EntityModel = lazy_loader.get_entity_model()
        if isinstance(entity_slug, str):
            entity_model = EntityModel.objects.for_user(
                user_model=user_model).get(
                slug__exact=entity_slug)
        elif isinstance(entity_slug, EntityModel):
            entity_model = entity_slug
        else:
            raise ValidationError('entity_slug must be an instance of str or EntityModel')

        LedgerModel = lazy_loader.get_ledger_model()
        acc_number = self.account_number
        ledger_model = LedgerModel.objects.create(
            entity=entity_model,
            posted=posted_ledger,
            # pylint: disable=unsubscriptable-object
            name=f'Bank Account {"***" + acc_number[-4:]}'
        )
        ledger_model.clean()
        self.ledger = ledger_model
        return self

    class Meta:
        abstract = True
        verbose_name = _('Bank Account')
        indexes = [
            models.Index(fields=['ledger']),
            models.Index(fields=['account_type']),
            models.Index(fields=['cash_account', 'account_type'])
        ]
        unique_together = [
            ('cash_account', 'account_number', 'routing_number')
        ]

    # pylint: disable=invalid-str-returned
    def __str__(self):
        # pylint: disable=bad-option-value
        return self.name

class BankAccountModel(BackAccountModelAbstract):
    """
    Base Bank Account Model
    """


import argparse

import os
import torch

import numpy as np
import skimage
import skimage.measure
import skimage.io
import shapely.geometry
import shapely.ops
from PIL import Image
from multiprocess import Pool
from tqdm import tqdm

from functools import partial

from lydorn_utils import print_utils, geo_utils

from frame_field_learning import polygonize_utils, plot_utils

DEBUG = False

def debug_print(s: str):
    if DEBUG:
        print_utils.print_debug(s)

def get_args():
    argparser = argparse.ArgumentParser(description=__doc__)
    argparser.add_argument(
        '--seg_filepath',
        required=True,
        nargs='*',
        type=str,
        help='Filepath(s) to input segmentation/mask image.')
    argparser.add_argument(
        '--im_dirpath',
        required=True,
        type=str,
        help='Path to the directory containing the corresponding images os the segmentation/mask. '
             'Files must have the same filename as --seg_filepath.'
             'Used for vizualization or saving the shapefile with the same coordinate system as that image.')
    argparser.add_argument(
        '--out_dirpath',
        required=True,
        type=str,
        help='Path to the output directory.')
    argparser.add_argument(
        '--out_ext',
        type=str,
        default="shp",
        choices=['pdf', 'shp'],
        help="File extension of the output. "
             "'pdf': pdf visualization (requires --im_dirpath for the image),  'shp': shapefile")
    argparser.add_argument(
        '--bbox',
        nargs='*',
        type=int,
        help='Selects area in bbox for computation.')

    args = argparser.parse_args()
    return args

def simplify(polygons, probs, tolerance):
    if type(tolerance) == list:
        out_polygons_dict = {}
        out_probs_dict = {}
        for tol in tolerance:
            out_polygons, out_probs = simplify(polygons, probs, tol)
            out_polygons_dict["tol_{}".format(tol)] = out_polygons
            out_probs_dict["tol_{}".format(tol)] = out_probs
        return out_polygons_dict, out_probs_dict
    else:
        out_polygons = [polygon.simplify(tolerance, preserve_topology=True) for polygon in polygons]
        return out_polygons, probs

def shapely_postprocess(out_contours, np_indicator, config):
    height = np_indicator.shape[0]
    width = np_indicator.shape[1]

    # Handle holes:
    line_string_list = [shapely.geometry.LineString(out_contour[:, ::-1]) for out_contour in out_contours]

    # Add image boundary line_strings for border polygons
    line_string_list.append(
        shapely.geometry.LinearRing([
            (0, 0),
            (0, height - 1),
            (width - 1, height - 1),
            (width - 1, 0),
        ]))

    # Merge polylines (for border polygons):
    multi_line_string = shapely.ops.unary_union(line_string_list)

    # Find polygons:
    polygons, dangles, cuts, invalids = shapely.ops.polygonize_full(multi_line_string)

    polygons = list(polygons)

    # Remove small polygons
    polygons = [polygon for polygon in polygons if
                config["min_area"] < polygon.area]

    # Remove low prob polygons
    filtered_polygons = []
    filtered_polygon_probs = []
    for polygon in polygons:
        prob = polygonize_utils.compute_geom_prob(polygon, np_indicator)
        # print("simple:", np_indicator.min(), np_indicator.mean(), np_indicator.max(), prob)
        if config["seg_threshold"] < prob:
            filtered_polygons.append(polygon)
            filtered_polygon_probs.append(prob)

    polygons, probs = simplify(filtered_polygons, filtered_polygon_probs, config["tolerance"])

    return polygons, probs

def polygonize(seg_batch, config, pool=None, pre_computed=None):
    # tic_total = time.time()

    assert len(seg_batch.shape) == 4 and seg_batch.shape[
        1] <= 3, "seg_batch should be (N, C, H, W) with C <= 3, not {}".format(seg_batch.shape)

    # Indicator
    # tic = time.time()
    indicator_batch = seg_batch[:, 0, :, :]
    np_indicator_batch = indicator_batch.cpu().numpy()
    # toc = time.time()
    # debug_print(f"Indicator to cpu: {toc - tic}s")

    if pre_computed is None or "init_contours_batch" not in pre_computed:
        # tic = time.time()
        init_contours_batch = polygonize_utils.compute_init_contours_batch(np_indicator_batch, config["data_level"], pool=pool)
        # toc = time.time()
        # debug_print(f"Init contours: {toc - tic}s")
    else:
        init_contours_batch = pre_computed["init_contours_batch"]

    # tic = time.time()
    # Convert contours to shapely polygons to handle holes:
    if pool is not None:
        shapely_postprocess_partial = partial(shapely_postprocess, config=config)
        polygons_probs_batch = pool.starmap(shapely_postprocess_partial, zip(init_contours_batch, np_indicator_batch))
        polygons_batch, probs_batch = zip(*polygons_probs_batch)
    else:
        polygons_batch = []
        probs_batch = []
        for i, out_contours in enumerate(init_contours_batch):
            polygons, probs = shapely_postprocess(out_contours, np_indicator_batch[i], config)
            polygons_batch.append(polygons)
            probs_batch.append(probs)

    # toc = time.time()
    # debug_print(f"Shapely post-process: {toc - tic}s")

    # toc_total = time.time()
    # debug_print(f"Total: {toc_total - tic_total}s")

    return polygons_batch, probs_batch

def run_one(seg_filepath, out_dirpath, config, im_dirpath, out_ext=None, bbox=None):
    filename = os.path.basename(seg_filepath)
    name = os.path.splitext(filename)[0]

    # Load image
    image = None
    im_filepath = os.path.join(im_dirpath, name + ".tif")
    if out_ext == "pdf":
        image = skimage.io.imread(im_filepath)

    # seg = skimage.io.imread(seg_filepath) / 255
    seg_img = Image.open(seg_filepath)
    seg = np.array(seg_img)
    if seg.dtype == np.uint8:
        seg = seg / 255
    elif seg.dtype == np.bool:
        seg = seg.astype(np.float)

    # Select bbox for dev
    if bbox is not None:
        assert len(bbox) == 4, "bbox should have 4 values"
        # bbox = [1440, 210, 1800, 650]  # vienna12
        # bbox = [2808, 2393, 3124, 2772]  # innsbruck19
        if image is not None:
            image = image[bbox[0]:bbox[2], bbox[1]:bbox[3]]
        seg = seg[bbox[0]:bbox[2], bbox[1]:bbox[3]]
        extra_name = ".bbox_{}_{}_{}_{}".format(*bbox)
    else:
        extra_name = ""

    # Convert to torch and add batch dim
    if len(seg.shape) < 3:
        seg = seg[:, :, None]
    seg_batch = torch.tensor(np.transpose(seg[:, :, :2], (2, 0, 1)), dtype=torch.float)[None, ...]

    out_contours_batch, out_probs_batch = polygonize(seg_batch, config)

    polygons = out_contours_batch[0]

    if out_ext == "shp":
        out_filepath = os.path.join(out_dirpath, name + ".shp")
        geo_utils.save_shapefile_from_shapely_polygons(polygons, im_filepath, out_filepath)
    elif out_ext == "pdf":
        base_filepath = os.path.splitext(seg_filepath)[0]
        filepath = base_filepath + extra_name + ".poly_simple.pdf"
        # plot_utils.save_poly_viz(image, polygons, filepath, linewidths=1, draw_vertices=True, color_choices=[[0, 1, 0, 1]])
        plot_utils.save_poly_viz(image, polygons, filepath, markersize=30, linewidths=1, draw_vertices=True)
    else:
        raise ValueError(f"out_ext should be shp or pdf, not {out_ext}")

def main():
    config = {
        "data_level": 0.5,
        "tolerance": 1.0,
        "seg_threshold": 0.5,
        "min_area": 10
    }
    # --- Process args --- #
    args = get_args()

    pool = Pool()
    list(tqdm(pool.imap(partial(run_one, out_dirpath=args.out_dirpath, config=config, im_dirpath=args.im_dirpath, out_ext=args.out_ext, bbox=args.bbox), args.seg_filepath), desc="Simple poly.", total=len(args.seg_filepath)))

if __name__ == '__main__':
    main()


from __future__ import annotations

from dataclasses import dataclass
import logging
from typing import Any, Dict, Generic, Iterable, Iterator, List, Optional, Set, TypeVar
from urllib.parse import urlparse

import attr
import requests
from tenacity.before_sleep import before_sleep_log

from swh.core.config import load_from_envvar
from swh.core.github.utils import GitHubSession
from swh.core.utils import grouper
from swh.scheduler import get_scheduler, model
from swh.scheduler.interface import SchedulerInterface

from . import USER_AGENT_TEMPLATE
from .utils import http_retry, is_valid_origin_url

logger = logging.getLogger(__name__)

@dataclass
class ListerStats:
    pages: int = 0
    origins: int = 0

    def __add__(self, other: ListerStats) -> ListerStats:
        return self.__class__(self.pages + other.pages, self.origins + other.origins)

    def __iadd__(self, other: ListerStats):
        self.pages += other.pages
        self.origins += other.origins

    def dict(self) -> Dict[str, int]:
        return {"pages": self.pages, "origins": self.origins}

StateType = TypeVar("StateType")
PageType = TypeVar("PageType")

BackendStateType = Dict[str, Any]
CredentialsType = Optional[Dict[str, Dict[str, List[Dict[str, str]]]]]

class Lister(Generic[StateType, PageType]):
    """The base class for a Software Heritage lister.

    A lister scrapes a page by page list of origins from an upstream (a forge, the API
    of a package manager, ...), and massages the results of that scrape into a list of
    origins that are recorded by the scheduler backend.

    The main loop of the lister, :meth:`run`, basically revolves around the
    :meth:`get_pages` iterator, which sets up the lister state, then yields the scrape
    results page by page. The :meth:`get_origins_from_page` method converts the pages
    into a list of :class:`model.ListedOrigin`, sent to the scheduler at every page. The
    :meth:`commit_page` method can be used to update the lister state after a page of
    origins has been recorded in the scheduler backend.

    The :func:`finalize` method is called at lister teardown (whether the run has
    been successful or not) to update the local :attr:`state` object before it's sent to
    the database. This method must set the :attr:`updated` attribute if an updated
    state needs to be sent to the scheduler backend. This method can call
    :func:`get_state_from_scheduler` to refresh and merge the lister state from the
    scheduler before it's finalized (and potentially minimize the risk of race
    conditions between concurrent runs of the lister).

    The state of the lister is serialized and deserialized from the dict stored in the
    scheduler backend, using the :meth:`state_from_dict` and :meth:`state_to_dict`
    methods.

    Args:
      scheduler: the instance of the Scheduler being used to register the
        origins listed by this lister
      url: a URL representing this lister, e.g. the API's base URL
      instance: the instance name, to uniquely identify this lister instance,
        if not provided the URL network location will be used
      credentials: dictionary of credentials for all listers. The first level
        identifies the :attr:`LISTER_NAME`, the second level the lister
        :attr:`instance`. The final level is a list of dicts containing the
        expected credentials for the given instance of that lister.
      max_pages: the maximum number of pages listed in a full listing operation
      max_origins_per_page: the maximum number of origins processed per page
      enable_origins: whether the created origins should be enabled or not

    Generic types:
      - *StateType*: concrete lister type; should usually be a :class:`dataclass` for
        stricter typing
      - *PageType*: type of scrape results; can usually be a :class:`requests.Response`,
        or a :class:`dict`

    """

    LISTER_NAME: str = ""
    github_session: Optional[GitHubSession] = None

    def __init__(
        self,
        scheduler: SchedulerInterface,
        url: str,
        instance: Optional[str] = None,
        credentials: CredentialsType = None,
        max_origins_per_page: Optional[int] = None,
        max_pages: Optional[int] = None,
        enable_origins: bool = True,
        with_github_session: bool = False,
    ):
        if not self.LISTER_NAME:
            raise ValueError("Must set the LISTER_NAME attribute on Lister classes")

        self.url = url
        if instance is not None:
            self.instance = instance
        else:
            self.instance = urlparse(url).netloc

        self.scheduler = scheduler

        if not credentials:
            credentials = {}
        self.credentials = list(
            credentials.get(self.LISTER_NAME, {}).get(self.instance, [])
        )

        # store the initial state of the lister
        self.state = self.get_state_from_scheduler()
        self.updated = False

        self.session = requests.Session()
        # Declare the USER_AGENT is more sysadm-friendly for the forge we list
        self.session.headers.update(
            {"User-Agent": USER_AGENT_TEMPLATE % self.LISTER_NAME}
        )
        self.github_session: Optional[GitHubSession] = (
            GitHubSession(
                credentials=credentials.get("github", {}).get("github", []),
                user_agent=str(self.session.headers["User-Agent"]),
            )
            if with_github_session
            else None
        )

        self.recorded_origins: Set[str] = set()
        self.max_pages = max_pages
        self.max_origins_per_page = max_origins_per_page
        self.enable_origins = enable_origins

    @http_retry(before_sleep=before_sleep_log(logger, logging.WARNING))
    def http_request(self, url: str, method="GET", **kwargs) -> requests.Response:

        logger.debug("Fetching URL %s with params %s", url, kwargs.get("params"))

        response = self.session.request(method, url, **kwargs)
        if response.status_code not in (200, 304):
            logger.warning(
                "Unexpected HTTP status code %s on %s: %s",
                response.status_code,
                response.url,
                response.content,
            )
        response.raise_for_status()

        return response

    def run(self) -> ListerStats:
        """Run the lister.

        Returns:
          A counter with the number of pages and origins seen for this run
          of the lister.

        """
        full_stats = ListerStats()
        self.recorded_origins = set()

        try:
            for page in self.get_pages():
                full_stats.pages += 1
                origins = list(self.get_origins_from_page(page))
                if (
                    self.max_origins_per_page
                    and len(origins) > self.max_origins_per_page
                ):
                    logger.info(
                        "Max origins per page set, truncated %s page results down to %s",
                        len(origins),
                        self.max_origins_per_page,
                    )
                    origins = origins[: self.max_origins_per_page]
                if not self.enable_origins:
                    logger.info(
                        "Disabling origins before sending them to the scheduler"
                    )
                    origins = [attr.evolve(origin, enabled=False) for origin in origins]
                sent_origins = self.send_origins(origins)
                self.recorded_origins.update(sent_origins)
                full_stats.origins = len(self.recorded_origins)
                self.commit_page(page)

                if self.max_pages and full_stats.pages >= self.max_pages:
                    logger.info("Reached page limit of %s, terminating", self.max_pages)
                    break
        finally:
            self.finalize()
            if self.updated:
                self.set_state_in_scheduler()

        return full_stats

    def get_state_from_scheduler(self) -> StateType:
        """Update the state in the current instance from the state in the scheduler backend.

        This updates :attr:`lister_obj`, and returns its (deserialized) current state,
        to allow for comparison with the local state.

        Returns:
          the state retrieved from the scheduler backend
        """
        self.lister_obj = self.scheduler.get_or_create_lister(
            name=self.LISTER_NAME, instance_name=self.instance
        )
        return self.state_from_dict(self.lister_obj.current_state)

    def set_state_in_scheduler(self) -> None:
        """Update the state in the scheduler backend from the state of the current
        instance.

        Raises:
          swh.scheduler.exc.StaleData: in case of a race condition between
            concurrent listers (from :meth:`swh.scheduler.Scheduler.update_lister`).
        """
        self.lister_obj.current_state = self.state_to_dict(self.state)
        self.lister_obj = self.scheduler.update_lister(self.lister_obj)

    # State management to/from the scheduler

    def state_from_dict(self, d: BackendStateType) -> StateType:
        """Convert the state stored in the scheduler backend (as a dict),
        to the concrete StateType for this lister."""
        raise NotImplementedError

    def state_to_dict(self, state: StateType) -> BackendStateType:
        """Convert the StateType for this lister to its serialization as dict for
        storage in the scheduler.

        Values must be JSON-compatible as that's what the backend database expects.
        """
        raise NotImplementedError

    def finalize(self) -> None:
        """Custom hook to finalize the lister state before returning from the main loop.

        This method must set :attr:`updated` if the lister has done some work.

        If relevant, this method can use :meth`get_state_from_scheduler` to merge the
        current lister state with the one from the scheduler backend, reducing the risk
        of race conditions if we're running concurrent listings.

        This method is called in a `finally` block, which means it will also run when
        the lister fails.

        """
        pass

    # Actual listing logic

    def get_pages(self) -> Iterator[PageType]:
        """Retrieve a list of pages of listed results. This is the main loop of the lister.

        Returns:
          an iterator of raw pages fetched from the platform currently being listed.
        """
        raise NotImplementedError

    def get_origins_from_page(self, page: PageType) -> Iterator[model.ListedOrigin]:
        """Extract a list of :class:`model.ListedOrigin` from a raw page of results.

        Args:
          page: a single page of results
        Returns:
          an iterator for the origins present on the given page of results
        """
        raise NotImplementedError

    def commit_page(self, page: PageType) -> None:
        """Custom hook called after the current page has been committed in the scheduler
        backend.

        This method can be used to update the state after a page of origins has been
        successfully recorded in the scheduler backend. If the new state should be
        recorded at the point the lister completes, the :attr:`updated` attribute must
        be set.

        """
        pass

    def send_origins(self, origins: Iterable[model.ListedOrigin]) -> List[str]:
        """Record a list of :class:`model.ListedOrigin` in the scheduler.

        Returns:
          the list of origin URLs recorded in scheduler database
        """
        valid_origins = []
        for origin in origins:
            if is_valid_origin_url(origin.url):
                valid_origins.append(origin)
            else:
                logger.warning("Skipping invalid origin: %s", origin.url)

        recorded_origins = []
        for batch_origins in grouper(valid_origins, n=1000):
            ret = self.scheduler.record_listed_origins(batch_origins)
            recorded_origins += [origin.url for origin in ret]

        return recorded_origins

    @classmethod
    def from_config(cls, scheduler: Dict[str, Any], **config: Any):
        """Instantiate a lister from a configuration dict.

        This is basically a backwards-compatibility shim for the CLI.

        Args:
          scheduler: instantiation config for the scheduler
          config: the configuration dict for the lister, with the following keys:
            - credentials (optional): credentials list for the scheduler
            - any other kwargs passed to the lister.

        Returns:
          the instantiated lister
        """
        # Drop the legacy config keys which aren't used for this generation of listers.
        for legacy_key in ("storage", "lister", "celery"):
            config.pop(legacy_key, None)

        # Instantiate the scheduler
        scheduler_instance = get_scheduler(**scheduler)

        return cls(scheduler=scheduler_instance, **config)

    @classmethod
    def from_configfile(cls, **kwargs: Any):
        """Instantiate a lister from the configuration loaded from the
        SWH_CONFIG_FILENAME envvar, with potential extra keyword arguments
        if their value is not None.

        Args:
            kwargs: kwargs passed to the lister instantiation
        """
        config = dict(load_from_envvar())
        config.update({k: v for k, v in kwargs.items() if v is not None})
        return cls.from_config(**config)

class StatelessLister(Lister[None, PageType], Generic[PageType]):
    def state_from_dict(self, d: BackendStateType) -> None:
        """Always return empty state"""
        return None

    def state_to_dict(self, state: None) -> BackendStateType:
        """Always set empty state"""
        return {}


import calendar
import datetime as dt
from pathlib import Path
from zipfile import ZipFile

import numpy as np

from src import functionsProspecAPI as prospec
from src.arrFiles import main as prep_files

def prep_n_run():
    name, path = model_params() # Recebe nome do modelo e caminho para os arquivos

    idServer = prospec.getIdOfServer("c5.24xlarge") # define o servidor que vai ser usado
    idQueue = prospec.getIdOfFirstQueueOfServer("c5.24xlarge") # define a fila que será usada

    display_studies() # mostra os arquivos criados desde então registrado no arquivo "estudos criados"

    studyId = int(
        input("Para qual estudo deseja enviar os arquivos?\n")) # recebe id do estudo

    prep_files() # baixa os gevazp e as ENAs do pluvia para o tratamento dos dados e os coloca nas pastas apropriadas

    # os decks devem ser inseridos manualmente na pasta de decks do estudo desejado
    # TODO #13 automatizar o download dos decks
    send_decks(path, studyId) # função que organiza e manda um zip com todos os decks necessários para rodar o estudo

    if name == "Curtíssimo prazo": # controle de fluxo para cada modelo

        todayDate, multipleRevision, targetDates = treat_dates() # recebe as datas tratadas para usar de parâmetro

        initialYear = todayDate.year # ano do parâmetro
        initialMonth = todayDate.month # mês do parâmetro
        duration = str(todayDate.month - targetDates[1].month)
        month = [targetDates[0].month, targetDates[1].month] # lista de parâmetros de mês
        year = [targetDates[0].year, targetDates[1].year] # lista de parâmetros de ano
        monthStr = '0' + \
            str(initialMonth) if initialMonth < 10 else str(initialMonth) #formatação da string de mês
        dateToFormat = (initialYear, monthStr)
        newaveFile = "NW%d%s.zip" % dateToFormat # nome do deck newave
        decompFile = "DC%d%s.zip" % dateToFormat # nome do deck decomp

        configFile = "Dados_Prospectivo.xlsx" # arquivo de configuração. obrigatoriamente colocado de forma manual na pasta do modelo

    prospec.sendFileToStudy(studyId, path+'/'+configFile, configFile) # manda arquivo de configuração para a plataforma

    path_prevs = path + "/prevs/" # cria a string de caminho para os prevs

    prospec.sendPrevsToStudy(studyId, path_prevs) # evia os prevs para a plataforma

    send_gevazp(path, studyId) # envia os arquivos do GEVAZP para a plataforma

    prospec.runExecution(studyId, idServer, idQueue, '', '0', '0', '2') # solicita a execução do estudo

# TODO #14 fazer o download dos resultados dos estudos da plataforma
    # prospec.generateNextRev(studyId, newaveFile, decompFile, configFile, [])

    # prospec.generateStudyDecks(studyId, initialYear, initialMonth, [duration], month, year, [False, False], [
    #                            True, True], newaveFile, [newaveFile, newaveFile], [decompFile, decompFile], [configFile, configFile], [])

def treat_dates(): # função que organiza as datas de acordo com a revisão do curtíssimo prazo
    targetDates = []
    multipleRevision = []
    todayDate = dt.date.today()
    firstRevDate = todayDate + dt.timedelta(weeks=1)
    secondRevDate = firstRevDate + dt.timedelta(weeks=1)
    daysDelta = 5 - int(firstRevDate.isoweekday())
    targetDates.append(firstRevDate - dt.timedelta(days=daysDelta))
    targetDates.append(secondRevDate - dt.timedelta(days=daysDelta))
    for target in targetDates:
        rev = get_rev(target)
        multipleRevision.append(rev)
    return todayDate, multipleRevision, targetDates

def get_rev(date_value): # calcula a revisão de interesse baseado num calendário
    day = date_value.day
    month = date_value.month
    year = date_value.year
    x = np.array(calendar.monthcalendar(year, month))
    revision = np.where(x == day)[0][0]
    return(revision)

def send_decks(path, uploadId): # envia arquivos de decks para a plataforma

    path_decks = path + "/Decks/" # cria caminho para a pasta de decks
    create_folder_structure(path_decks) # cria as pastas com os conteúdos no formato que o zip precisa ter
    completeStudyZip = "CompleteStudy.zip" # nome do zip final
    zip_path = path_decks+completeStudyZip # criação do caminho para o zip na pasta

    # TODO #11 criação do zip apropriado para upload
    # os caminhos dentro do zip se referem à pasta raiz do projeto. Eles têm que se referir à pasta study
    # do modelo (e.g. CP/Curtíssimo/study/<pasta>/<arquivo> deveria ser somente /<pasta>/<arquivo>)
    decks = [item for item in Path(path_decks).glob(
        "**/*") if item.suffix == ".zip"] #lista os arquivos zip da pasta de decks

    create_zip(zip_path, path_decks) # função que cria o zip a ser enviado

    prospec.sendFileToStudy(uploadId, zip_path, completeStudyZip) # envia o zip para o estudo de interesse
    zipDecks = Path(zip_path) # cria objeto de caminho
    prospec.completeStudyDecks(uploadId, zipDecks.name, []) # solicita o tratamento do zip na plataforma com os decks completos
    zipDecks.unlink() # remove zip enviado para evitar um zip recursivo

def create_zip(zip_path, path_decks):
    studyFiles = Path(path_decks+'/study') # cria endereço para os arquivos do zip

    # with ZipFile(zip_path, 'w') as zp:
    #     for deck in decks:
    #         if "decomp" in deck.parts:
    #             arcName = "/DECOMP/%s" % deck.name

    #         else:
    #             arcName = "/NEWAVE/%s" % deck.name
    #         zp.write(deck, arcName)

def create_folder_structure(path_decks):
    extract_dc(path_decks) # extrai os arquivos do decomp para a criação de pastas individuais para eles
    folders(path_decks) # extrai os arquivos pertinentes aos decks a serem enviados
    rm_tmp(path_decks) # limpa arquivos temporários

def folders(path_decks):
    files = [item for item in Path(path_decks).glob("**/*") if item.suffix == ".zip"] # lista os zips da pasta de decks
    for file in files:
        if ("sem" in file.name or "NW" in file.name) and not ("Relatorio" in file.name): # acha os zips que não são de relatório
            with ZipFile(file, 'r') as zp:
                zp.extractall(path_decks+"/study/"+file.stem) # extrai para a pasta do zip

def extract_dc(path_decks):
    for file in Path(path_decks).glob("**/*"): # itera arquivos da pasta de decks
        if file.suffix == ".zip" and "DC" in file.name: # condicional para pegar somente os arquivos do DECOMP
            with ZipFile(file) as zp:
                zp.extractall(path_decks+'/tmp') # extração para uma pasta temporária

def rm_tmp(path_decks): # remoção da pasta temporária
    tmp = Path(path_decks+"/tmp")
    for file in tmp.glob('**/*'):
        file.unlink()
    tmp.rmdir()

def send_gevazp(path, uploadId): # envio dos arquivos do GEVAZP para o estudo
    path_gevazp = path + "/GEVAZP/"
    for file in Path(path_gevazp).glob("**/*"):
        if file.is_file():
            prospec.sendFileToDeck(uploadId, "", file, file.name)

def studiesInfo(): # retorna informações de um estudo
    display_studies() # disponibiliza o histórico de estudos criados
    infoId = input("Qual estudo deseja?\n") # recebe id do estudo
    infoStudy = prospec.getInfoFromStudy(infoId) # solicita à plataforma as infos do estudo
    if infoStudy == None: # tratamento de resposta de um estudo que não está mais presente na plataforma
        print("Esse estudo não está mais na plataforma!")
    else:
        print(infoStudy)

def create_study(nameStudy): # cria um estudo
    idStudy = prospec.createStudy(nameStudy, "", 0, 0) # cria um estudo na plataforma com o nome dado bno parâmetro

    with open("estudos criados", 'a') as fp: # registra no histórico a criação do estudo
        fp.write("ID: %s, Nome: %s\n" % (idStudy, nameStudy))

    print("O estudo %s foi criado com ID %s" % (nameStudy, idStudy))

def display_studies(): # imprime na tela o histórico de estudos
    print("Esses são os estudos criados até então:")
    with open("estudos criados", 'r') as fp:
        for line in fp:
            print(line)

def model_params(): # escolhe o nome e caminho do estudo
    nameStudy = "" # recebe o nome apropriado ou continua vazio em caso de erro
    path = "" # mesmo que o nome
    choice = input(
        "Qual o tipo de estudo que deseja?\n1- Curtísimo prazo\n2- ONS CP\n3- Matriz CP\n") # recebe o tipo de estudo a ser feito
    choice = int(choice)
    if choice == 1:
        nameStudy = "Curtíssimo prazo" # define o nome do estudo
        path = "CP/Curtissimo" # define o caminho para os arquivos do estudo

    elif choice < 4:
        # TODO #15 Controle de nomes e caminhos dos demais estudos

        print("Opção em implementação.")

    else:
        print("Opção inválida.")

    return nameStudy, path

def main():
    prospec.authenticateProspec(
        "daniel.mazucanti@skoposenergia.com.br", "Skopos2020") # segurança em primeiro lugar, não é mesmo?

    numRequests = prospec.getNumberOfRequests() # recebe o número de requisições feitas

    print("Foram feitas %s requisições até o momento." % numRequests) # imprime o número de requisições feitas

    while True: # cria laço para impressão do menu
        control_flow = input(
            '1- Criar estudo\n2- Rodar estudo\n3- Abortar execução\n4- Informações do estudo\n0- Sair\n')
            # recebe opção do menu
        if control_flow == "1": # criar estudo

            nameStudy, path_opt = model_params() # recebe os parâmetros do estudo desejado

            create_study(nameStudy) # cria estudo na plataforma

        elif control_flow == "2": # rodar estudo

            prep_n_run() # prepara os arquivos, os envia para a plataforma e solicita a execução

        elif control_flow == "3": # pelo direito de escolher
            display_studies() # mostra histórico de estudos criados
            stopId = int(input("Qual estudo deseja parar?\n")) # recebe id do estudo a ser abortado
            prospec.abortExecution(stopId) # solicita o interrupção do estudo na plataforma

        elif control_flow == "4": # infos do estudo
            studiesInfo() # duh

        else: # qualquer coisa que não esteja nos casos anteriores para aqui
            print("Programa encerrado.")
            break

# create_folder_structure("CP/Curtissimo/Decks/")
# isso aqui era para debugar essa função


from operator import index
import pandas as pd
import joblib
import os
import argparse
import numpy as np

from tqdm import tqdm
from sklearn.preprocessing import LabelBinarizer

def main(input_path, output_path):
    all_paths = os.listdir(input_path)
    folder_paths = [folder_path for folder_path in all_paths if os.path.isdir(os.path.join(input_path, folder_path))]
    print(f"Folder paths: {folder_paths}")
    print(f"Number of folders :{len(folder_paths)}")

    # create_labels = ["basketball", "boxing", "chess"]

    data = pd.DataFrame()

    image_formats = ['jpg', 'JPG', 'PNG', 'png']
    labels = []
    counter = 0
    for i, folder_path in tqdm(enumerate(folder_paths), total=len(folder_paths)):
        # if folder_path not in create_labels:
        #     continue
        image_paths = os.listdir(os.path.join(input_path, folder_path))
        label = folder_path

        for image_path in image_paths:
            if image_path.split('.')[-1] in image_formats:
                data.loc[counter, 'image_path'] = f"{input_path}/{folder_path}/{image_path}"
                labels.append(label)
                counter += 1

    labels = np.array(labels)
    lb = LabelBinarizer()
    labels = lb.fit_transform(labels)

    print(labels.shape)
    if len(labels[0]) == 1:
        for i in range(len(labels)):
            index = labels[i]
            data.loc[i, 'target'] = int(index)
    elif len(labels[0]) > 1:
        for i in range(len(labels)):
            index = np.argmax(labels[i])
            data.loc[i, 'target'] = int(index)

    # shuffle the dataset
    data = data.sample(frac=1).reset_index(drop=True)

    print(f"Number of classes: {len(lb.classes_)}")
    print(f"Total instances: {len(data)}")

    # save to csv
    data.to_csv(os.path.join(input_path,"data.csv"), index=False)

    # saving the binarized labels
    print("Saving the binarized labels in pickled file")
    joblib.dump(lb, os.path.join(output_path, "lb.pkl"))

    print(data.head(5))

if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("-i", "--input_path", required=True, help="Path to the root folder containing dataset images")
    ap.add_argument("-o", "--output_path", required=True, help="Path to save the preprocessed data")

    args = ap.parse_args()
    arguments = args.__dict__

    main(**arguments)


# -*- coding: utf-8 -*-
import os
import threading
import tkinter as tk # imports
from tkinter import ttk
# from performance_function_test.write_performance_data import getPerformance_data
from tkinter import messagebox
import time
import yaml
import inspect
import ctypes

from openpyxl import load_workbook

# from performance_function_test.write_performance_data import write_performance_data
from performance_function_test.write_performance_data import creat_excel, write_data_to_excel

import sys
sys.path.append("..")
from file_util import *

global projectid,path_report
projectid=1
path_report=get_base_path() + 'performance_function_test\\report'
path_yaml = get_base_path() + 'performance_function_test\\yaml'

def init_tab2(tab_page):
    global text_monkey_info
    monty = ttk.LabelFrame(tab_page, text=' ')
    monty.grid(column=0, row=0, padx=20, pady=20)

    tk.Button(monty, text='点击配置文件 ', font=("Arial", 11), width=30, height=1,command=lambda: openReport(path_yaml)).grid(column=0, row=0, sticky='W', padx=20, pady=10)
    tk.Button(monty, text='点击开始运行 ', font=("Arial", 11), width=30, height=1,command=run_function_test).grid(column=0, row=5, sticky='W', padx=20,
                                                                                pady=10)
    tk.Button(monty, text='查看测试结果 ', font=("Arial", 11), width=30, height=1,command=lambda: openReport(path_report)).grid(column=0, row=9, sticky='W', padx=20,
                                                                                pady=10)

def get_function():
    log_path = get_base_path() + 'performance_function_test\\log\\log.txt'
    path_log = os.path.abspath(log_path)
    excel_path = creat_excel(path_yaml)

    with open(get_base_path() + 'performance_function_test\\yaml\\performance_caps.yaml', 'r', encoding='utf-8') as file:
        data = yaml.load(file)
    scene_info = data['scene']
    # print(scene_info)
    process_info = data['process']
    lens_process = len(process_info)
    lens_scene = len(scene_info)

    averageValueLog_path = get_base_path() + 'performance_function_test\log\\averageValueLog.txt'
    path_averageValueLog = os.path.abspath(averageValueLog_path)
    file_averageValueLog= open(path_averageValueLog, 'w')

    isCustomStop = False # 是否用户取消
    all_times_total = [0] * len(scene_info) # 保存每次的平均值总和，每个item是一个场景的平均值列表
    has_value_times = [0] * len(scene_info) # 保存每个场景有数据的次数是多少次
    all_times_average = [0] * len(scene_info)
    for n in range(1,4):
        # ask = messagebox.askquestion('APP性能测试', '是否开始第 '+ str(n)+' 轮测试（测试会有3轮测试）')
        # print(ask)
        # while askquestion == 'yes':

        y =2
        scene_index =0
        for temp_scene in scene_info.values():
            # f = open(path_log, 'w')
            # f.closed
            f = open(path_log, 'w')

            start_info = tk.messagebox.askquestion(title='第'+ str(n) +'轮测试：' + temp_scene,message='请对页面进行操作，操作完成后，点击“确定”，进行数据采集')
            print(start_info)
            th = threading.Thread(target=set_text_test, args=(f,))
            if start_info=='yes':
                print('数据采集中')
                th.setDaemon(True)  # 守护线程
                th.start()
            else:
                isCustomStop = True
                break
            end_info=tk.messagebox.showinfo('数据采集中', '正在采集数据，请至少采集30S的数据，\n点击" 确定 "可结束数据采集')

            # tk.Label(monty,text='路径为：', font=("Arial", 11), width=20, height=2).grid(column=1,row=6, sticky='W')
            if end_info =='ok':
                _async_raise(th.ident,SystemExit)
                print('关闭txt文件')
                f.close()
                averageValue=write_data_to_excel(path_log, excel_path, 4, y, '第' + str(n) + '次测试')

                if len(averageValue)>0:
                    has_value_times[scene_index] = has_value_times[scene_index] + 1

                tmp = all_times_total[scene_index]
                if type(tmp).__name__ == 'list' and len(tmp)==len(averageValue):
                    for m in range(len(tmp)):
                        tmp[m] =tmp[m] + averageValue[m]
                    all_times_total[scene_index] = tmp
                else:
                    all_times_total[scene_index]=averageValue

            print('写入成功')
            scene_index = scene_index +1
            y = y + lens_process*2 + 2

            with open('结果存放.txt', 'a') as file_handle:  # .txt可以不自己新建,代码会自动新建
                file_handle.write('\n')  # 有时放在循环里面需要自动转行，不然会覆盖上一条数据
        # print(file=file_averageValueLog, end='\n')
        if isCustomStop:
            break
    if not isCustomStop:
        tk.messagebox.showinfo('APP性能测试', '测试已完成，请查看测试结果')

        # 求总平均值
        for v1 in range(len(all_times_total)):
            tmp = [0] * len(all_times_total[v1])

            for v2 in range(len(all_times_total[v1])):
                if has_value_times[v1] == 0:
                    break
                tmp[v2] = all_times_total[v1][v2]/has_value_times[v1]
            all_times_average[v1] =tmp
        print("has_value_times:",has_value_times)
        print("all_times_average:",all_times_average)

        print('将平均值写入excel')

        workbook = load_workbook(excel_path)
        workbook.guess_types = True  # 猜测格式类型
        sheet_average = workbook.get_sheet_by_name('APP性能测试')

        begin_row = 2
        begin_col = 2

        for i in range(len(all_times_average)):
            if (len(all_times_average[i]) == 0):
                continue
            for j in range(len(all_times_average[i])):
                sheet_average.cell(row=begin_row + j, column=begin_col + i, value=round(all_times_average[i][j], 2))

        workbook.save(excel_path)
        workbook.close()

    # file_averageValueLog.close()
# #运行多线程
# def get_function(project_id,insert_text):
#     getPerformance_data(project_id,insert_text)

def set_text_test(f):
    global get_adb_kugou_isrun
    get_adb_kugou_isrun = True

    while get_adb_kugou_isrun:
        get_adb_info(f)

def get_adb_kugou_stop(log_path,excel_path,n,y):
    global get_adb_kugou_isrun
    get_adb_kugou_isrun = False
    write_data_to_excel(log_path, excel_path, 4, y, '第' + str(n) + '次测试')
    # insert_text( '第 ' + str(n) + ' 次测试__' + '场景 ' + str(y) + ' 测试结束了\n')

def run_function_test():
    th = threading.Thread(target=get_function())
    th.setDaemon(True)  # 守护线程
    th.start()

#往txt中插入数据
def insert_text(message):
    if(message==None):
        message = ""

    text_monkey_info.insert(tk.END,message)

def openReport(path_report):

    path = os.path.abspath(path_report)  # 获取当前脚本所在的路径
    os.startfile(path)

def _async_raise(tid, exctype):
   """raises the exception, performs cleanup if needed"""
   tid = ctypes.c_long(tid)
   if not inspect.isclass(exctype):
      exctype = type(exctype)
   res = ctypes.pythonapi.PyThreadState_SetAsyncExc(tid, ctypes.py_object(exctype))
   if res == 0:
      raise ValueError("invalid thread id")
   elif res != 1:
      # """if it returns a number greater than one, you're in trouble,
      # and you should call it again with exc=NULL to revert the effect"""
      ctypes.pythonapi.PyThreadState_SetAsyncExc(tid, None)
      raise SystemError("PyThreadState_SetAsyncExc failed")


from flask import Flask, render_template
import os
app = Flask(__name__)

@app.route('/<img>')
def hello_world(img):
    print(img)
    images_list = os.listdir('./static/images')
    print (images_list)
    return render_template('hello.html', images = images_list, selected_image = img)

@app.route('/')
def hello_world_2():
    images_list = os.listdir('./static/images')
    print (images_list)
    return render_template('hello.html', images = images_list, selected_image = None)

if __name__ =="__main__":
    app.run(host="0.0.0.0")


"""Version string parser and generator."""

import collections.abc
import enum
import itertools
import logging
import re
import typing as t

import packaging.version
import semver

_LOG = logging.getLogger(__name__)

PY_PRE_RELEASE_INDICATORS = {'a', 'b', 'c', 'rc'}

@enum.unique
class VersionComponent(enum.IntEnum):
    """Enumeration of standard version components."""

    Major = 1 << 1
    Minor = 1 << 2
    Patch = 1 << 3
    Release = Major | Minor | Patch
    PrePatch = 1 << 4
    DevPatch = 1 << 5
    Local = 1 << 6

def _version_tuple_checker(version_tuple, flags):
    return all([(_ is not None if flag else _ is None) for _, flag in zip(version_tuple, flags)])

class Version(collections.abc.Hashable):  # pylint: disable = too-many-public-methods
    """For storing and manipulating version information.

    Definitions of acceptable version formats are provided in readme.
    """

    _re_number = r'(?:0|[123456789][0123456789]*)'
    # _re_sha = r'[0123456789abcdef]+'
    _re_letters = r'(?:[abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ]+)'
    _pattern_letters = re.compile(_re_letters)
    _re_alphanumeric = r'(?:[0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ]+)'
    _pattern_alphanumeric = re.compile(_re_alphanumeric)
    _re_sep = r'(?:[\.-])'

    _re_release_parts = \
        r'(?P<major>{n})(?:\.(?P<minor>{n}))?(?:\.(?P<patch>{n}))?'.format(n=_re_number)
    _pattern_release = re.compile(_re_release_parts)

    @classmethod
    def _parse_release_str(cls, release: str) -> tuple:
        match = cls._pattern_release.fullmatch(release)
        assert match is not None
        major_match = match.group('major')
        assert major_match is not None
        major = int(major_match)
        minor_match = match.group('minor')
        if minor_match is not None:
            minor = int(minor_match)
        else:
            minor = None
        patch_match = match.group('patch')
        if patch_match is not None:
            patch = int(patch_match)
        else:
            patch = None
        return major, minor, patch

    _re_pre_separator = rf'(?P<preseparator>{_re_sep})'
    _re_pre_type = rf'(?P<pretype>{_re_letters})'
    _re_pre_patch = rf'(?P<prepatch>{_re_number})'
    _re_pre_release_part = rf'{_re_pre_separator}?{_re_pre_type}?{_re_pre_patch}?'
    _pattern_pre_release_part = re.compile(_re_pre_release_part)
    _re_pre_release_parts = r'(?:{0}{2})|(?:{0}?{1}{2}?)'.format(_re_sep, _re_letters, _re_number)
    _pattern_pre_release = re.compile(_re_pre_release_parts)
    _pattern_pre_release_check = re.compile(rf'(?:{_re_pre_release_parts})+')

    @classmethod
    def _parse_pre_release_str(cls, pre_release: str) -> t.Sequence[
            t.Tuple[t.Optional[str], t.Optional[str], t.Optional[int]]]:
        parts = cls._pattern_pre_release.findall(pre_release)
        _LOG.debug('parsed pre-release string %s into %s',
                   repr(pre_release), parts)
        tuples = []
        for part in parts:
            match = cls._pattern_pre_release_part.fullmatch(part)
            assert match is not None
            pre_patch_match = match.group('prepatch')
            if pre_patch_match is not None:
                pre_patch = int(pre_patch_match)
            else:
                pre_patch = None
            tuples.append((match.group('preseparator'), match.group('pretype'), pre_patch))
        return tuples

    _re_local_separator = rf'({_re_sep})'
    _re_local_part = rf'({_re_alphanumeric})'
    _re_local_parts = rf'\+{_re_local_part}(?:{_re_local_separator}{_re_local_part})*'
    _pattern_local = re.compile(_re_local_parts)

    @classmethod
    def _parse_local_str(cls, local: str) -> tuple:
        match = cls._pattern_local.fullmatch(local)
        assert match is not None
        return tuple([_ for _ in match.groups() if _ is not None])

    _re_release = r'(?P<release>{n}(?:\.{n})?(?:\.{n})?)'.format(n=_re_number)
    _re_pre_release = r'(?P<prerelease>(?:(?:{0}{2})|(?:{0}?{1}{2}?))+)'.format(
        _re_sep, _re_letters, _re_number)
    _re_local = r'(?P<local>\+{0}([\.-]{0})*)'.format(_re_alphanumeric)
    # _re_named_parts_count = 3 + 3
    _re_version = rf'{_re_release}{_re_pre_release}?{_re_local}?'
    _pattern_version = re.compile(_re_version)

    @classmethod
    def from_str(cls, version_str: str):
        """Create version from string."""
        match = cls._pattern_version.fullmatch(version_str)  # type: t.Optional[t.Match[str]]
        if match is None:
            raise ValueError(f'version string {repr(version_str)} is invalid')
        _LOG.debug('version_query parsed version string %s into %s: %s %s',
                   repr(version_str), type(match), match.groupdict(), match.groups())

        _release = match.group('release')
        _pre_release = match.group('prerelease')
        _local = match.group('local')

        major, minor, patch = cls._parse_release_str(_release)
        pre_release = None if _pre_release is None else cls._parse_pre_release_str(_pre_release)
        local = None if _local is None else cls._parse_local_str(_local)

        return cls(major=major, minor=minor, patch=patch, pre_release=pre_release, local=local)

    @classmethod
    def from_tuple(cls, version_tuple: tuple):
        return cls(*version_tuple)

    @classmethod
    def from_dict(cls, version_dict: dict):
        return cls(**version_dict)

    @classmethod
    def from_py_version(cls, py_version: packaging.version.Version):
        """Create version from a standard Python version object."""
        if not isinstance(py_version, packaging.version.Version):
            _LOG.warning('attempting to parse %s as packaging.version.Version...', type(py_version))
        ver = py_version._version
        major, minor, patch = [ver.release[i] if len(ver.release) > i
                               else None for i in range(3)]
        pre_release: t.Optional[t.Sequence[
            t.Tuple[t.Optional[str], t.Optional[str], t.Optional[int]]]] = None
        local = None
        pre_ver: t.Optional[t.Tuple[None, int]]
        if len(ver.release) == 4:
            pre_ver = (None, ver.release[3])
        elif len(ver.release) > 4:
            raise NotImplementedError(py_version)
        else:
            pre_ver = None
        pre_ver_present = sum(1 for _ in (ver.post, ver.dev, ver.pre) if _)
        if pre_ver and pre_ver_present:
            raise NotImplementedError(py_version)
        if pre_ver_present > 1:
            raise NotImplementedError(py_version)
        if ver.dev:
            pre_ver = ver.dev
        elif ver.pre:
            pre_ver = ver.pre
        if pre_ver:
            prefix_dot = None if pre_ver[0] in PY_PRE_RELEASE_INDICATORS else '.'
            pre_release = [(prefix_dot, pre_ver[0], pre_ver[1] if len(pre_ver) > 1 else None)]
        if ver.local:
            local = tuple(itertools.chain.from_iterable(
                (dot, str(_)) for dot, _ in zip('.' * len(ver.local), ver.local)))[1:]
        _LOG.debug('parsing %s %s', type(py_version), py_version)
        return cls(major, minor, patch, pre_release=pre_release, local=local)

    @classmethod
    def from_sem_version(cls, sem_version: t.Union[dict, semver.VersionInfo]):
        """Create version from semantic version object."""
        _LOG.debug('parsing %s %s', type(sem_version), sem_version)
        if isinstance(sem_version, semver.VersionInfo):
            major, minor, patch = sem_version.major, sem_version.minor, sem_version.patch
            pre_release = sem_version.prerelease
            local = sem_version.build
        else:
            major, minor, patch = sem_version['major'], sem_version['minor'], sem_version['patch']
            pre_release = sem_version['prerelease']
            local = sem_version['build']
        if pre_release is not None:
            raise NotImplementedError(sem_version)
        if local is not None:
            local = cls._parse_local_str(f'+{local}')
        return cls(major, minor, patch, pre_release=pre_release, local=local)

    @classmethod
    def from_version(cls, version: 'Version'):
        return cls.from_dict(version.to_dict())

    def __init__(
            self, major: int, minor: t.Optional[int] = None, patch: t.Optional[int] = None, *args,
            pre_release: t.Sequence[
                t.Tuple[t.Optional[str], t.Optional[str], t.Optional[int]]] = None,
            local: t.Union[str, tuple] = None):
        self._major: t.Optional[int] = None
        self._minor: t.Optional[int] = None
        self._patch: t.Optional[int] = None
        self._pre_release: t.Optional[
            t.List[t.Tuple[t.Optional[str], t.Optional[str], t.Optional[int]]]] = None
        self._local: t.Optional[t.Sequence[str]] = None

        self.release = major, minor, patch

        if args and pre_release is not None and local is not None:
            raise ValueError(f'args={args}, pre_release={pre_release} and local={local}'
                             f' are all present in {repr(self)}')

        if pre_release is None:
            pre_release, consumed_args = self._get_pre_release_from_args(args)
            if consumed_args > 0:
                args = args[consumed_args:]
            else:
                pre_release = None
        self.pre_release = pre_release

        if args and local is not None:
            raise ValueError(f'args={args} and local={local} are present at the same time'
                             f' in {repr(self)}')

        if local is None:
            if len(args) == 1 and isinstance(args[0], tuple):
                local = args[0]
            elif not args:
                local = None
            else:
                local = args
        elif isinstance(local, str):
            local = (local,)
        self.local = local

    def _get_pre_release_from_args(self, args) -> t.Tuple[
            t.Sequence[t.Tuple[t.Optional[str], t.Optional[str], t.Optional[int]]], int]:
        pre_release: t.List[t.Tuple[t.Optional[str], t.Optional[str], t.Optional[int]]] = []
        consumed_args = 0
        if args and isinstance(args[0], tuple):
            for i, arg in enumerate(args):
                if not isinstance(arg, tuple):
                    break
                if len(arg) == 3 and arg[0] in (None, '.', '-'):
                    pre_release.append(arg)
                    consumed_args += 1
                    continue
                if i == len(args) - 1:
                    break
                raise ValueError(f'pre-release segment arg={arg} (index {i} in args={args}'
                                 f' in {repr(self)}) must be a 3-tuple')
        else:
            accumulated: t.List[t.Union[int, str]] = []
            for i, arg in enumerate(args):
                if not accumulated:
                    if arg in (None, '.', '-'):
                        if len(args) < i + 3:
                            raise ValueError(f'expected 3 consecutive values from index {i}'
                                             f' in args={args} in {repr(self)}')
                    else:
                        break
                accumulated.append(arg)
                consumed_args += 1
                if len(accumulated) == 3:
                    pre_release.append(tuple(accumulated))
                    accumulated = []
        return pre_release, consumed_args

    @property
    def release(self) -> t.Tuple[int, t.Optional[int], t.Optional[int]]:
        assert self._major is not None
        return self._major, self._minor, self._patch

    @release.setter
    def release(self, release: t.Tuple[int, t.Optional[int], t.Optional[int]]):
        # major: int, minor: t.Optional[int] = None, patch: t.Optional[int] = None):
        if not isinstance(release, tuple):
            raise TypeError(
                f'release={repr(release)} is of wrong type {type(release)} in {repr(self)}')
        if len(release) != 3:
            raise ValueError(
                f'release={repr(release)} has wrong length {len(release)} in {repr(self)}')

        major, minor, patch = release

        if not isinstance(major, int):
            raise TypeError(f'major={repr(major)} is of wrong type {type(major)} in {repr(self)}')
        if major < 0:
            raise ValueError(f'major={repr(major)} has wrong value in {repr(self)}')
        if minor is not None and not isinstance(minor, int):
            raise TypeError(f'minor={repr(minor)} is of wrong type {type(minor)} in {repr(self)}')
        if minor is not None and minor < 0:
            raise ValueError(f'minor={repr(minor)} has wrong value in {repr(self)}')
        if patch is not None and not isinstance(patch, int):
            raise TypeError(f'patch={repr(patch)} is of wrong type {type(patch)} in {repr(self)}')
        if patch is not None and patch < 0:
            raise ValueError(f'patch={repr(patch)} has wrong value in {repr(self)}')
        if minor is None and patch is not None:
            raise ValueError(f'patch={repr(patch)} is present but not minor in {repr(self)}')

        self._major = major
        self._minor = minor
        self._patch = patch

    @property
    def pre_release(self) -> t.Optional[
            t.Sequence[t.Tuple[t.Optional[str], t.Optional[str], t.Optional[int]]]]:
        """Pre-release version component."""
        if self._pre_release is None:
            return None
        return self._pre_release.copy()

    @pre_release.setter
    def pre_release(
            self, pre_release: t.Optional[
                t.List[t.Tuple[t.Optional[str], t.Optional[str], t.Optional[int]]]]):
        if pre_release is None:
            self._pre_release = None
            return

        if not isinstance(pre_release, collections.abc.Sequence):
            raise TypeError(f'pre_release={repr(pre_release)} is of wrong type {type(pre_release)}'
                            f' in {repr(self)}')
        if len(pre_release) == 0:
            raise ValueError(f'pre_release has no elements although it is set in {repr(self)}')

        for pre in pre_release:
            if not isinstance(pre, tuple):
                raise TypeError(
                    f'pre-release part {repr(pre)} is of wrong type {type(pre)} in {repr(self)}')
            if len(pre) != 3:
                raise ValueError(
                    f'pre-release part={repr(pre)} has wrong length {len(pre)} in {repr(self)}')
            pre_separator, pre_type, pre_patch = pre
            self._check_pre_release_parts(pre_separator, pre_type, pre_patch)

        self._pre_release = pre_release

    def _check_pre_release_parts(self, pre_separator, pre_type, pre_patch):
        """Verify that the given pre-release version identifier parts are valid."""
        if pre_separator is not None and not isinstance(pre_separator, str):
            raise TypeError(f'pre_separator={repr(pre_separator)} is of wrong type'
                            f' {type(pre_separator)} in {repr(self)}')
        if pre_separator is not None and pre_separator not in ('-', '.'):
            raise ValueError(f'pre_separator={repr(pre_separator)} has wrong value in {repr(self)}')
        if pre_type is not None and not isinstance(pre_type, str):
            raise TypeError(
                f'pre_type={repr(pre_type)} is of wrong type {type(pre_type)} in {repr(self)}')
        if pre_type is not None and type(self)._pattern_letters.fullmatch(pre_type) is None:
            raise ValueError(f'pre_type={repr(pre_type)} has wrong value in {repr(self)}')
        if pre_patch is not None and not isinstance(pre_patch, int):
            raise TypeError(
                f'pre_patch={repr(pre_patch)} is of wrong type {type(pre_patch)} in {repr(self)}')
        if pre_patch is not None and pre_patch < 0:
            raise ValueError(f'pre_patch={repr(pre_patch)} has wrong value in {repr(self)}')
        if pre_separator is None and pre_type is None and pre_patch is not None:
            raise ValueError(f'neither pre_separator nor pre_type is set'
                             f' but pre_patch={repr(pre_patch)} is in {repr(self)}')
        if pre_separator is not None and pre_type is None and pre_patch is None:
            raise ValueError(f'pre_separator={repr(pre_separator)} is present'
                             f' but neither pre_type nor pre_patch is in {repr(self)}')

    @property
    def has_pre_release(self):
        return self._pre_release is not None

    @property
    def local(self) -> t.Optional[t.Sequence[str]]:
        return self._local

    @local.setter
    def local(self, local: t.Optional[t.Sequence[str]]):
        if local is None:
            self._local = None
            return

        if not isinstance(local, collections.abc.Sequence):
            raise TypeError(f'local={repr(local)} is of wrong type {type(local)} in {repr(self)}')

        if len(local) % 2 != 1:
            raise ValueError(f'local={repr(local)} has wrong length {len(local)} in {repr(self)}')

        for i, part in enumerate(local):
            if not isinstance(part, str):
                raise TypeError(f'local_part or local_separator {repr(part)} is of wrong type'
                                f' {type(part)} in {repr(self)}')
            if i % 2 == 0:
                if type(self)._pattern_alphanumeric.fullmatch(part) is None:
                    raise ValueError(f'local_part={repr(part)} has wrong value in {repr(self)}')
            elif part not in ('-', '.'):
                raise ValueError(f'local_separator={repr(part)} has wrong value in {repr(self)}')

        self._local = tuple(local)

    @property
    def has_local(self):
        return self._local is not None

    def increment(self, component: VersionComponent, amount: int = 1) -> 'Version':
        """Increment a selected version component and return self."""
        if not isinstance(component, VersionComponent):
            raise TypeError(f'component={repr(component)} is of wrong type {type(component)}')
        if not isinstance(amount, int):
            raise TypeError(f'amount={repr(amount)} is of wrong type {type(amount)}')
        if amount < 1:
            raise ValueError(f'amount={amount} has wrong value')

        if component in (VersionComponent.Major, VersionComponent.Minor, VersionComponent.Patch):
            self._increment_release(component, amount)

        elif component is VersionComponent.PrePatch:
            self._increment_pre_path(amount)

        elif component is VersionComponent.DevPatch:
            if self._pre_release is None:
                self._pre_release = []
            if not self._pre_release or self.pre_release_to_tuple(True)[-1][1] != 'dev':
                self._pre_release.append(('.', 'dev', amount))
            else:
                pre_sep, pre_type, pre_patch = self._pre_release[-1]
                if pre_patch is None:
                    pre_patch = amount
                else:
                    pre_patch += amount
                self._pre_release[-1] = (pre_sep, pre_type, pre_patch)
        else:
            raise ValueError(f'incrementing component={repr(component)} is not possible')

        return self

    def _increment_release(self, component: VersionComponent, amount: int):
        if component in (VersionComponent.Major, VersionComponent.Minor):
            if component is VersionComponent.Major:
                assert self._major is not None
                self._major += amount
                if self._minor is not None:
                    self._minor = 0
            else:
                if self._minor is None:
                    self._minor = amount
                else:
                    self._minor += amount

            if self._patch is not None:
                self._patch = 0

        else:
            if self._minor is None:
                self._minor = 0
            if self._patch is None:
                self._patch = amount
            else:
                self._patch += amount

        self._pre_release = None
        self._local = None

    def _increment_pre_path(self, amount):
        if self._pre_release is None or self.pre_release_to_tuple(True)[0][1] != '':
            self._pre_release = [('-', None, amount)]
        else:
            pre_sep, pre_type, pre_patch = self._pre_release[0]
            assert isinstance(pre_patch, int), (type(pre_patch), pre_patch)
            pre_patch += amount
            self.pre_release = [(pre_sep, pre_type, pre_patch)]

    def devel_increment(self, new_commits: int = 1) -> 'Version':
        """Increment version depending on current version number and return self.

        If there's no dev patch, then increment release patch by one and increment dev patch
        by number of new commits.

        If there's a dev patch, then increment it by number of new commits.
        """
        if not self.has_pre_release or self.pre_release_to_tuple(True)[-1][1] != 'dev':
            self.increment(VersionComponent.Patch)
        self.increment(VersionComponent.DevPatch, new_commits)

        return self

    def release_to_str(self) -> str:
        """Get string representation of this version's release component."""
        version_tuple = self._major, self._minor, self._patch
        if _version_tuple_checker(version_tuple, (True, False, False)):
            return '.'.join(str(_) for _ in version_tuple[:1])
        if _version_tuple_checker(version_tuple, (True, True, False)):
            return '.'.join(str(_) for _ in version_tuple[:2])
        if _version_tuple_checker(version_tuple, (True, True, True)):
            return '.'.join(str(_) for _ in version_tuple[:3])
        raise ValueError(f'cannot generate valid version string from {repr(self)}')

    def _pre_release_segment_to_str(self, segment: int) -> str:
        assert self._pre_release is not None
        version_tuple = self._pre_release[segment]
        if _version_tuple_checker(version_tuple, (True, True, False)):
            return '{}{}'.format(*version_tuple[:2])
        if _version_tuple_checker(version_tuple, (True, False, True)):
            return f'{version_tuple[0]}{version_tuple[2]}'
        if _version_tuple_checker(version_tuple, (False, True, True)):
            return '{}{}'.format(*version_tuple[1:])
        if _version_tuple_checker(version_tuple, (True, True, True)):
            return '{}{}{}'.format(*version_tuple)
        raise ValueError(f'cannot generate valid version string from {repr(self)}')

    def pre_release_to_str(self) -> str:
        if self._pre_release is None:
            return ''
        return ''.join(self._pre_release_segment_to_str(i)
                       for i, _ in enumerate(self._pre_release))

    def local_to_str(self) -> str:
        if not self._local:
            return ''
        return f'+{"".join(self._local)}'

    def to_str(self) -> str:
        return f'{self.release_to_str()}{self.pre_release_to_str()}{self.local_to_str()}'

    def release_to_tuple(self, sort: bool = False) -> tuple:
        return (0 if sort else None) if self._major is None else self._major, \
            (0 if sort else None) if self._minor is None else self._minor, \
            (0 if sort else None) if self._patch is None else self._patch

    def pre_release_segment_to_tuple(self, segment: int, sort: bool = False) -> tuple:
        assert self._pre_release is not None
        pre_separator, pre_type, pre_patch = self._pre_release[segment]
        return (1 if pre_type is None else 0) if sort else pre_separator, \
            ('' if pre_type is None else pre_type.lower()) if sort else pre_type, \
            (0 if sort else None) if pre_patch is None else pre_patch

    def pre_release_to_tuple(self, sort: bool = False) -> tuple:
        """Create tuple from this version's pre-release component."""
        if self._pre_release is None:
            if sort:
                return ((1, '', 0),)
            return ()
        parts = [self.pre_release_segment_to_tuple(i, sort)
                 for i, _ in enumerate(self._pre_release)]
        return tuple(parts) if sort else tuple(itertools.chain.from_iterable(parts))

    def local_to_tuple(self, sort: bool = False) -> tuple:
        if self._local is None:
            return ()
        return tuple(0 if _ in ('.', '-') else _.lower() for _ in self._local) \
            if sort else self._local

    def to_tuple(self, sort: bool = False) -> tuple:
        return self.release_to_tuple(sort) + self.pre_release_to_tuple(sort) \
            + self.local_to_tuple(sort)

    def to_dict(self) -> dict:
        return {field[1:]: value for field, value in vars(self).items()}

    def to_py_version(self) -> packaging.version.Version:
        return packaging.version.Version(self.to_str())

    def to_sem_version(self) -> semver.VersionInfo:
        return semver.VersionInfo.parse(self.to_str())

    def __repr__(self):
        fields = ', '.join(f'{field[1:]}: {repr(value)}' for field, value in vars(self).items())
        return f'{type(self).__name__}({fields})'

    def __str__(self):
        return self.to_str()

    def __hash__(self):
        return hash(self.to_tuple(sort=True))

    def __eq__(self, other):
        return not self < other and not other < self

    def __ne__(self, other):
        return self < other or other < self

    def __gt__(self, other):
        return other < self

    def __ge__(self, other):
        return not self < other

    def __lt__(self, other):
        if not isinstance(other, Version):
            raise TypeError(f'cannot compare {type(self)} and {type(other)}')

        self_release = self.release_to_tuple(True)
        other_release = other.release_to_tuple(True)
        if self_release != other_release:
            return self_release < other_release

        self_pre_release = self.pre_release_to_tuple(True)
        other_pre_release = other.pre_release_to_tuple(True)
        for self_part, other_part in itertools.zip_longest(
                self_pre_release, other_pre_release, fillvalue=(1, '', 0)):
            if self_part != other_part:
                return self_part < other_part

        self_local = self.local_to_tuple(True)
        other_local = other.local_to_tuple(True)
        for self_part, other_part in zip(self_local, other_local):
            if self_part != other_part:
                return self_part < other_part
        if len(self_local) != len(other_local):
            return len(self_local) < len(other_local)

        return False

    def __le__(self, other):
        return not other < self


import gymnasium as gym
from ray.rllib.core.rl_module.rl_module import SingleAgentRLModuleSpec
from torchvision.models import resnet18, ResNet18_Weights
from ray.rllib.core.rl_module.torch.torch_rl_module import TorchRLModule
from ray.rllib.algorithms.ppo import PPOConfig
from ray.rllib.utils.nested_dict import NestedDict
from ray.rllib.core.testing.torch.bc_module import DiscreteBCTorchModule
from ray.rllib.algorithms.ppo.ppo_catalog import PPOCatalog
from ray.rllib.core.models.specs.typing import SpecType
from ray.rllib.core.rl_module.rl_module import RLModuleConfig
from ray.rllib.algorithms.ppo.ppo_catalog import PPOCatalog
from typing import Mapping, Any
import cv2
import numpy as np
from ray.rllib.models.torch.torch_distributions import TorchCategorical
from pprint import pprint
from ray.rllib.policy.sample_batch import SampleBatch
#env = gym.make("ALE/BeamRider-v5")
import torch
import torch.nn as nn
import torch.utils.data
import torch.nn.functional as F

from typing import Mapping, Any

from ray.rllib.algorithms.ppo.ppo_rl_module import PPORLModule

from ray.rllib.core.models.base import ACTOR, CRITIC, ENCODER_OUT, STATE_IN
from ray.rllib.core.rl_module.rl_module import RLModule
from ray.rllib.core.rl_module.torch import TorchRLModule
from ray.rllib.policy.sample_batch import SampleBatch
from ray.rllib.utils.annotations import override
from ray.rllib.utils.framework import try_import_torch
from ray.rllib.utils.nested_dict import NestedDict

torch, nn = try_import_torch()

class ResDown(nn.Module):
    """
    Residual down sampling block for the encoder
    """
    def __init__(self, channel_in, channel_out, scale=2):
        super(ResDown, self).__init__()

        self.conv1 = nn.Conv2d(channel_in, channel_out//2, 3, 1, 1)
        self.BN1 = nn.BatchNorm2d(channel_out//2)
        self.conv2 = nn.Conv2d(channel_out//2, channel_out, 3, 1, 1)
        self.BN2 = nn.BatchNorm2d(channel_out)

        self.conv3 = nn.Conv2d(channel_in, channel_out, 3, 1, 1)

        self.AvePool = nn.AvgPool2d(scale, scale)

    def forward(self, x):
        skip = self.conv3(self.AvePool(x))

        x = F.rrelu(self.BN1(self.conv1(x)))
        x = self.AvePool(x)
        x = self.BN2(self.conv2(x))

        x = F.rrelu(x + skip)
        return x

class ResUp(nn.Module):
    """
    Residual up sampling block for the decoder
    """
    def __init__(self, channel_in, channel_out, scale=2):
        super(ResUp, self).__init__()

        self.conv1 = nn.Conv2d(channel_in, channel_out//2, 3, 1, 1)
        self.BN1 = nn.BatchNorm2d(channel_out//2)
        self.conv2 = nn.Conv2d(channel_out//2, channel_out, 3, 1, 1)
        self.BN2 = nn.BatchNorm2d(channel_out)

        self.conv3 = nn.Conv2d(channel_in, channel_out, 3, 1, 1)

        self.UpNN = nn.Upsample(scale_factor = scale,mode = "nearest")

    def forward(self, x):
        skip = self.conv3(self.UpNN(x))

        x = F.rrelu(self.BN1(self.conv1(x)))
        x = self.UpNN(x)
        x = self.BN2(self.conv2(x))

        x = F.rrelu(x + skip)
        return x

class Encoder(nn.Module):
    """
    Encoder block
    Built for a 3x64x64 image and will result in a latent vector of size z x 1 x 1
    As the network is fully convolutional it will work for images LARGER than 64
    For images sized 64 * n where n is a power of 2, (1, 2, 4, 8 etc) the latent feature map size will be z x n x n

    When in .eval() the Encoder will not sample from the distribution and will instead output mu as the encoding vector
    and log_var will be None
    """
    def __init__(self, channels, ch=64, z=512):
        super(Encoder, self).__init__()
        self.conv1 = ResDown(channels, ch)  # 64
        self.conv2 = ResDown(ch, 2*ch)  # 32
        self.conv3 = ResDown(2*ch, 4*ch)  # 16
        self.conv4 = ResDown(4*ch, 8*ch)  # 8
        self.conv5 = ResDown(8*ch, 8*ch)  # 4
        self.conv_mu = nn.Conv2d(8*ch, z, 2, 2)  # 2
        self.conv_log_var = nn.Conv2d(8*ch, z, 2, 2)  # 2

    def sample(self, mu, log_var):
        std = torch.exp(0.5*log_var)
        eps = torch.randn_like(std)
        return mu + eps*std

    def forward(self, x):
        x = x.float()
        x = x.permute(0, 3, 1, 2)
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        x = self.conv4(x)
        x = self.conv5(x)
        #x = x.flatten(start_dim=1)

        if self.training:
            mu = self.conv_mu(x)
            log_var = self.conv_log_var(x)
            x = self.sample(mu, log_var)
        else:
            mu = self.conv_mu(x)
            x = mu
            log_var = None

        return x, mu, log_var

class Decoder(nn.Module):
    """
    Decoder block
    Built to be a mirror of the encoder block
    """

    def __init__(self, channels, ch=64, z=512):
        super(Decoder, self).__init__()
        self.conv1 = ResUp(z, ch*8)
        self.conv2 = ResUp(ch*8, ch*8)
        self.conv3 = ResUp(ch*8, ch*4)
        self.conv4 = ResUp(ch*4, ch*2)
        self.conv5 = ResUp(ch*2, ch)
        self.conv6 = ResUp(ch, ch//2)
        self.conv7 = nn.Conv2d(ch//2, channels, 3, 1, 1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        x = self.conv4(x)
        x = self.conv5(x)
        x = self.conv6(x)
        x = self.conv7(x)

        return x

class VAE(nn.Module):
    """
    VAE network, uses the above encoder and decoder blocks
    """
    def __init__(self, channel_in, ch=64, z=512):
        super(VAE, self).__init__()
        """Res VAE Network
        channel_in  = number of channels of the image
        z = the number of channels of the latent representation
        (for a 64x64 image this is the size of the latent vector)
        """

        self.encoder = Encoder(channel_in, ch=ch, z=z)
        self.decoder = Decoder(channel_in, ch=ch, z=z)

    def forward(self, x):
        mu = self.encoder(x)[1]

        #encoding, mu, log_var = self.encoder(x)
        #recon = self.decoder(encoding)
        #HARDCODED
        #return recon, mu, log_var
        #return mu
        return torch.flatten(mu, start_dim=1)

class CusEnv(gym.Env):
    def __init__(self,env_config):
        self.env = gym.make("ALE/BeamRider-v5", full_action_space=True)
        self.action_space = self.env.action_space
        self.observation_space = gym.spaces.Box(0, 255, (84, 84, 3), np.uint8) #self.env.observation_space

    def reset(self, seed=None, options=None):
        obs,info = self.env.reset()
        return cv2.resize(obs, (84, 84)), info

    def step(self, action):
        obs, reward, terminated, truncated, info = self.env.step(action)
        obs = cv2.resize(obs, (84, 84))
        return obs, reward, terminated, truncated, info

class PPOTorchRLModule(PPORLModule, TorchRLModule):
    framework: str = "torch"

    def __init__(self, *args, **kwargs):
        TorchRLModule.__init__(self, *args, **kwargs)
        PPORLModule.__init__(self, *args, **kwargs)
        #self.blahencoder = Encoder(channels=3).cuda()
        self.blahencoder = VAE(channel_in=3, ch=64)
        #checkpoint = torch.load("/lab/kiran/ckpts/pretrained/atari/" + "STL10_ATTARI_64.pt")
        #self.blahencoder.load_state_dict(checkpoint['model_state_dict'])
        #self._weights = ResNet18_Weights.IMAGENET1K_V1
        #self.blahencoder = resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)
        self.blahencoder.eval()
        for param in self.blahencoder.parameters():
            param.requires_grad = False
        #self._preprocess = self._weights.transforms()
        #print(self.encoder)

    def _forward_inference(self, batch: NestedDict) -> Mapping[str, Any]:
        output = {}

        # TODO (Artur): Remove this once Policy supports RNN
        if self.encoder.config.shared:
            batch[STATE_IN] = None
        else:
            batch[STATE_IN] = {
                ACTOR: None,
                CRITIC: None,
            }
        batch[SampleBatch.SEQ_LENS] = None

        with torch.no_grad():
            batch['obs'] = self.blahencoder(batch['obs'].cuda()).detach()

        #batch['obs'] = self.blahencoder(batch['obs'].cuda())
        encoder_outs = self.encoder(batchself.vf(encoder_outs[ENCODER_OUT][CRITIC]))
        # TODO (Artur): Un-uncomment once Policy supports RNN
        # output[STATE_OUT] = encoder_outs[STATE_OUT]

        # Actions
        action_logits = self.pi(encoder_outs[ENCODER_OUT][ACTOR])
        output[SampleBatch.ACTION_DIST_INPUTS] = action_logits

        return output

    def _forward_exploration(self, batch: NestedDict) -> Mapping[str, Any]:
        """PPO forward pass during exploration.
        Besides the action distribution, this method also returns the parameters of the
        policy distribution to be used for computing KL divergence between the old
        policy and the new policy during training.
         """
        with torch.no_grad():
            batch['obs'] = self.blahencoder(batch['obs'].cuda()).detach()
            return self._common_forward(batch)

    def _forward_train(self, batch: NestedDict) -> Mapping[str, Any]:
        with torch.no_grad():
            batch['obs'] = self.blahencoder(batch['obs'].cuda()).detach()
        return self._common_forward(batch)

    def _common_forward(self, batch: NestedDict) -> Mapping[str, Any]:
        output = {}

        # TODO (Artur): Remove this once Policy supports RNN
        if self.encoder.config.shared:
            batch[STATE_IN] = None
        else:
            batch[STATE_IN] = {
                ACTOR: None,
                CRITIC: None,
            }
        batch[SampleBatch.SEQ_LENS] = None

        #batch['obs'] = self.blahencoder(batch['obs'].cuda())
        encoder_outs = self.encoder(batch)
        # TODO (Artur): Un-uncomment once Policy supports RNN
        # output[STATE_OUT] = encoder_outs[STATE_OUT]

        # Value head
        vf_out = self.vf(encoder_outs[ENCODER_OUT][CRITIC])
        output[SampleBatch.VF_PREDS] = vf_out.squeeze(-1)

        # Policy head
        action_logits = self.pi(encoder_outs[ENCODER_OUT][ACTOR])
        output[SampleBatch.ACTION_DIST_INPUTS] = action_logits

        return output

spec = SingleAgentRLModuleSpec(
    module_class=PPOTorchRLModule,
    observation_space=gym.spaces.Box(-1, 1, (512,)),
    action_space=gym.spaces.Discrete(18),
    model_config_dict={"vf_share_layers": True,
        "fcnet_hiddens": [256, 256]}
)

config = (
    PPOConfig()
    .environment(CusEnv, env_config={
        "frameskip": 1,
        "full_action_space": False,
        "repeat_action_probability": 0.0
    })
    .training(_enable_learner_api=True)
    .rl_module(
        _enable_rl_module_api=True,
        rl_module_spec=spec,
    )
    .training(model={"vf_share_layers": True,
        "fcnet_hiddens": [256, 256],
        "conv_filters": [[16, 4, 2], [32, 4, 2], [64, 4, 2], [128, 4, 2]]},
        lambda_=0.95,
        lr=0.0001,
        kl_coeff=0.5,
        clip_param=0.1,
        vf_clip_param=10.0,
        grad_clip=100.0,
        entropy_coeff=0.01,
        train_batch_size=5000,
        sgd_minibatch_size=500,
        num_sgd_iter=10)
    #.exploration(exploration_config={})
    .reporting(min_time_s_per_iteration=30)
    .rollouts(num_rollout_workers=46,
        num_envs_per_worker=1,
        rollout_fragment_length='auto')
    .resources(
        num_gpus_per_learner_worker = .4,
        num_gpus=0,
        num_learner_workers=1,
        num_cpus_per_worker=0.2,
        num_gpus_per_worker=.14)
        #num_cpus_for_local_worker = 1,
)

algorithm = config.build()

# run for 2 training steps
for _ in range(2000000):
    result = algorithm.train()
    pprint(result)


from collections import deque
from dataclasses import dataclass
import random
from typing import Deque

from loguru import logger
import networkx

@dataclass
class OnlyRandomWalkerStrategyParams:
    length_of_memory: int = 1
    pass

class OnlyRandomWalker:
    """
    Dieser Random Walker macht zufällige Schritte und hat ein Gedächtnis von
    Knoten, welche er versucht zu vermeiden. Er erstellt
    niemals neue Kanten und verändert den Graphen nicht.
    """

    name: str
    g: networkx.Graph
    searched_information: int
    start_point: tuple

    length_of_memory: int

    last_nodes: Deque[tuple]

    @classmethod
    def withParams(
        cls,
        name: str,
        g: networkx.Graph,
        searched_information: int,
        start_point: tuple,
        random_walker_params: OnlyRandomWalkerStrategyParams,
    ):
        return cls(name, g, searched_information, start_point, random_walker_params.length_of_memory)

    def __init__(
        self, name: str, g: networkx, searched_information: int,  start_point: tuple, length_of_memory: int
    ):
        self.name = name
        self.g = g
        self.searched_information = searched_information
        self.start_point = start_point

        self.length_of_memory = length_of_memory

        self.last_nodes = deque([], length_of_memory)

        self.last_nodes.appendleft(start_point)

        logger.trace(f"RandomWalker {self.name} started at {self.current_node}")

    def current_node(self):
        return self.last_nodes[0]

    def walk(self) -> list[tuple]:
        # Don't visit nodes in history
        possible_nodes = list(self.g.neighbors(self.current_node()))
        possible_nodes = [
            node for node in possible_nodes if node not in self.last_nodes
        ]

        # Wenn wir keine Nachbarn haben, die nicht schon besucht wurden,
        # dann müssen wir halt einen besuchen, den wir schon besucht haben.
        if len(possible_nodes) == 0:
            possible_nodes = list(self.g.neighbors(self.current_node()))

        next_node = random.choice(possible_nodes)

        self.last_nodes.appendleft(next_node)

        # Kein Rückgabewert, dieser Random Walker gibt niemals
        # eine neue Kante zurück. Er dient nur zum Vergleichen.
        return []