from pyspark.sql import SparkSession
from pyspark.sql.functions import from_json, to_json, col, unbase64, base64, split, expr
from pyspark.sql.types import StructField, StructType, StringType, BooleanType, ArrayType, DateType

redisMessageSchema = StructType(
    [
        StructField("key", StringType()),
        StructField("value", StringType()),
        StructField("expiredType", StringType()),
        StructField("expiredValue",StringType()),
        StructField("existType", StringType()),
        StructField("ch", StringType()),
        StructField("incr",BooleanType()),
        StructField("zSetEntries", ArrayType( \
            StructType([
                StructField("element", StringType()),\
                StructField("score", StringType())   \
            ]))                                      \
        )

    ]
)

# TO-DO: create a StructType for the Payment schema for the following fields:
# {"reservationId":"9856743232","customerName":"Frank Aristotle","date":"Sep 29, 2020, 10:06:23 AM","amount":"946.88"}

payment_schema = StructType([
    StructField("reservationId", StringType()),
    StructField("customerName", StringType()),
    StructField("date", StringType()),
    StructField("amount", StringType())
])

# TO-DO: create a spark session, with an appropriately named application name
# TO-DO: set the log level to WARN

spark_session = SparkSession.builder.appName("redis-payment").getOrCreate()
spark_session.sparkContext.setLogLevel("WARN")

#TO-DO: read the redis-server kafka topic as a source into a streaming dataframe with the bootstrap server localhost:9092, configuring the stream to read the earliest messages possible
#TO-DO: using a select expression on the streaming dataframe, cast the key and the value columns from kafka as strings, and then select them
#TO-DO: using the redisMessageSchema StructType, deserialize the JSON from the streaming dataframe
# TO-DO: create a temporary streaming view called "RedisData" based on the streaming dataframe
# it can later be queried with spark.sql
#TO-DO: using spark.sql, select key, zSetEntries[0].element as payment from RedisData
#TO-DO: from the dataframe use the unbase64 function to select a column called payment with the base64 decoded JSON, and cast it to a string
#TO-DO: using the payment StructType, deserialize the JSON from the streaming dataframe, selecting column customer.* as a temporary view called Customer
#TO-DO: using spark.sql select reservationId, amount,  from Payment

redis_raw_stream_df = spark_session\
   .readStream\
   .format("kafka")\
   .option("kafka.bootstrap.servers","localhost:9092")\
   .option("subscribe","redis-server")\
   .option("startingOffsets","earliest")\
   .load()

redis_stream_df = redis_raw_stream_df.selectExpr("CAST(key AS STRING) key", "CAST(value AS STRING) value")

redis_stream_df\
    .withColumn("value", from_json("value", redisMessageSchema))\
    .select(col("value.*"))\
    .createOrReplaceTempView("redis_vw")

payment_stream_df = spark_session.sql("select key, zSetEntries[0].element as payment from redis_vw")

payment_stream_df\
    .withColumn("payment", unbase64(payment_stream_df.payment).cast("string"))\
    .withColumn("payment", from_json("payment", payment_schema))\
    .select(col("payment.*"))\
    .createOrReplaceTempView("payment_vw")

new_payment_stream_df = spark_session.sql("select reservationId, amount from payment_vw where amount is not null")

# TO-DO: write the stream in JSON format to a kafka topic called payment-json, and configure it to run indefinitely, the console output will not show any output.
#You will need to type /data/kafka/bin/kafka-console-consumer --bootstrap-server localhost:9092 --topic customer-attributes --from-beginning to see the JSON data
#
# The data will look like this: {"reservationId":"9856743232","amount":"946.88"}

new_payment_stream_df\
    .selectExpr("cast(reservationId as string) as key", "to_json(struct(*)) as value")\
    .writeStream\
    .format("kafka")\
    .option("kafka.bootstrap.servers","localhost:9092")\
    .option("topic","payment")\
    .option("checkpointLocation","/tmp/kafkacheckpoint")\
    .start()\
    .awaitTermination()

# to run
# /home/workspace/spark/bin/spark-submit --packages org.apache.spark:spark-sql-kafka-0-10_2.11:2.4.6 /home/workspace/payment-json-fields.py | tee /home/workspace/spark/logs/payment.log

# to check the data in kafka
# kafka-console-consumer --topic payment --bootstrap-server localhost:9092 --from-beginning


import asyncio
from functools import wraps

from quart import (
    Quart, render_template,
    request, websocket
)

app = Quart(__name__)

connected_websockets = set()

def collect_websocket(func):
    @wraps(func)
    async def wrapper(*args, **kwargs):
        global connected_websockets
        queue = asyncio.Queue()
        connected_websockets.add(queue)
        try:
            return await func(queue, *args, **kwargs)
        finally:
            connected_websockets.remove(queue)
    return wrapper

async def broadcast(message):
    global connected_websockets
    for queue in connected_websockets:
        await queue.put(message)

@app.route('/')
async def index():
    return await render_template('main.html')

@app.websocket('/ws')
@collect_websocket
async def ws(queue):
    await websocket.accept()
    while True:
        data = await queue.get()
        await websocket.send_json(data)

@app.route('/telepath', methods=['POST'])
async def telepath():
    data = await request.get_json()
    await broadcast(data.get('message', ''))
    return {}

if __name__ == '__main__':
    app.run(port=5000)


import hashlib
import inspect
import json
from typing import Any

def fn_args() -> dict[str, Any]:
    """Gets the arguments of the function that called this function."""
    frame = inspect.currentframe()
    assert frame is not None and frame.f_back is not None
    keys, _, _, local = inspect.getargvalues(frame.f_back)
    args = {k: local[k] for k in keys}
    args.update(local.get("kwargs", {}))
    args.update(local.get("_kwargs", {}))
    return args

def human_size(num_bytes: int, binary: bool = True) -> str:
    def float_fmt(f: float) -> str:
        return f"{f:.2f}".rstrip("0").rstrip(".")

    units = ["B", "KiB", "MiB", "GiB"]
    base = 1024
    if not binary:
        units = ["B", "KB", "MB", "GB"]
        base = 1000

    v = float(num_bytes)
    for unit in units[:-1]:
        if abs(v) < base:
            return f"{float_fmt(v)} {unit}"
        v /= base
    return f"{float_fmt(v)} {units[-1]}"

def stable_hash(val: Any) -> int:
    val = json.dumps(
        val,
        ensure_ascii=False,
        sort_keys=True,
        indent=None,
        separators=(",", ":"),
    )
    val = hashlib.md5(val.encode("utf-8")).digest()
    return int.from_bytes(val, "big")


#!/usr/bin/env python3

import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
import tkinter as tk
import numpy as np
import numpy.linalg
import argparse
import yaml

class LTISystem:
    """Class representing a LTI system given in state-space representation where A := system matrix, B := input matrix, C := output matrix, D := feedtrough matrix."""

    def __init__(self, A, B, C, D):
        """Initializes a LTI system."""

        self.__taylor_coeff = [1.0,
                               1.0,
                               0.5,
                               0.16666666666,
                               0.04166666666,
                               0.00833333333,
                               0.00138888888,
                               0.00019841269,
                               0.00002480158,
                               0.00000275573,
                               0.00000027557]
        x_dim = A.shape[1]
        u_dim = B.shape[1]
        assert(A.shape[0] == x_dim) # domain-dim has to be state-dim
        self.__A = A
        assert(B.shape[0] == x_dim) # domain-dim has to be state-dim
        self.__B = B
        assert(C.shape[1] == x_dim) # codomain-dim has to be state-dim
        self.__C = C
        assert(D.shape[0] == u_dim) # domain-dim has to be control-dim
        assert(D.shape[1] == x_dim) # codomain-dim has to be state-dim
        self.__D = D

    def discretize(self, delta_t, taylor_order=10):
        """Discretizes system matrix and intput matrix up to the 10th taylor order."""

        A_d = np.zeros(self.__A.shape)
        discretization_factor_B = np.zeros(self.__A.shape)
        delta_t_pow = 1.0
        A_pow = np.identity(self.__A.shape[0])
        for i in range(0, taylor_order+1):
            A_d = A_d + (self.__taylor_coeff[i]*A_pow*delta_t_pow)
            delta_t_pow = delta_t*delta_t_pow
            if (i < taylor_order):
                discretization_factor_B = discretization_factor_B + (self.__taylor_coeff[i+1]*A_pow*delta_t_pow)
            A_pow = self.__A*A_pow
        B_d = discretization_factor_B*self.__B

        return A_d, B_d

    def calculate_trajectory(self, A_d, B_d, x, u, num_steps):
        """Calculates the phase space state trajectory given an initial state and control input."""

        x_trajectory = x
        x_i = x
        for i in range(0, num_steps):
            x_i = (A_d*x_i) + (B_d*u)
            x_trajectory = np.hstack((x_trajectory, x_i))

        return x_trajectory

    def calculate_equilibrium(self, u):
        """Calculates the LTI system equilibrium. Returns 'None' if there is no equilibrium or infinitely many."""

        x_eq = None
        try:
            x_eq = np.linalg.solve(self.__A, -self.__B*u)
        except numpy.linalg.LinAlgError:
            print("There is no equilibrium or infinitely many!")

        return x_eq

    def calculate_equilibrium_as_function_of_u(self, u_bounds, u_ticks):
        """Calculates the LTI system equilibrium as a function of the control input. Returns 'None' if there is at least once no equilibrium or infinitely many."""

        init = False
        for u in np.arange(u_bounds[0, 0], u_bounds[0, 1], (u_bounds[0, 1]-u_bounds[0, 0])/u_ticks):
            x_eq = self.calculate_equilibrium(u)
            if x_eq is None:
                return None
            if init is False:
                x_eq_of_u = x_eq
                init = True
            else:
                x_eq_of_u = np.hstack((x_eq_of_u, x_eq))

        return x_eq_of_u

    def calculate_transfer_function_poles(self):
        """Calculates the LTI system transfer function."""

        # TODO recheck: does this only hold when system is controllable and observable?
        w, v = np.linalg.eig(self.__A)

        return w

    def generate_vector_field(self, u, xi, xj):
        """Generates the phase space vector field over two state elements."""

        K_u = self.__B*u
        xi_dot = self.__A[0, 0]*xi+self.__A[0, 1]*xj+K_u[0, 0]
        xj_dot = self.__A[1, 0]*xi+self.__A[1, 1]*xj+K_u[1, 0]

        return xi_dot, xj_dot

class BicycleModel(LTISystem):
    """Class representing the bicycle model."""

    def __init__(self, m, v, theta, c_f, c_r, l_f, l_r):
        """Initializes the bicycle model from its system parameters."""

        A = np.matrix([[-(c_f+c_r)/(m*v), -(m*v*v+(c_r*l_r-c_f*l_f))/(m*v*v)],
                      [(c_r*l_r-c_f*l_f)/theta, -(c_r*l_r*l_r+c_f*l_f*l_f)/(theta*v)]])
        B = np.matrix([[c_f/(m*v)], [(c_f*l_f)/theta]])
        C = np.matrix([[1.0, 0.0]])
        D = np.matrix([[0.0, 0.0]])
        LTISystem.__init__(self, A, B, C, D)

class BicycleModelApp(tk.Frame):
    def __init__(self, master, config):
        tk.Frame.__init__(self, master)
        self.master = master

        # set title
        master.title("Bicycle Model")

        # init system parameters
        self.__m = config["m_init"]
        self.__v = config["v_init"]
        self.__theta = config["theta_init"]
        self.__c_f = config["c_f_init"]
        self.__c_r = config["c_r_init"]
        self.__l_f = config["l_f_init"]
        self.__l_r = config["l_r_init"]

        # init state and control input
        self.__x = np.matrix([[config["slip_angle_init"]], [config["yaw_rate_init"]]])
        self.__u = np.matrix([[config["steering_angle_init"]]])

        # init number of prediction steps
        self.__num_steps = int(config["num_steps_init"])
        self.__delta_t = config["delta_t"]

        # init other parameters
        self.__x_bounds = np.matrix([config["slip_angle_bounds"], config["yaw_rate_bounds"]])
        self.__u_bounds = np.matrix([config["steering_angle_bounds"]])
        self.__x_ticks = int(config["x_ticks"])
        self.__u_ticks = int(config["u_ticks"])

        # init figure
        self.__f = plt.Figure(figsize=(15, 5))

        # calculate data for the plots
        self.__xi = None
        self.__xj = None
        self.__xi_dot = None
        self.__xj_dot = None
        self.__time = None
        self.__x_eq = None
        self.__x_eq_of_u = None
        self.__x_trajectory = None
        self.__poles = None
        self.recalculate_plot_data()

        # draw plots on figure
        self.__phase_space_plot = None
        self.__trajectory_plot = None
        self.__trajectory_plot_start = None
        self.__trajectory_plot_end = None
        self.__trajectory_plot_0 = None
        self.__trajectory_plot_1 = None
        self.__laplace_0 = None
        self.__laplace_1 = None
        self.__x_eq_plot = None
        self.__x_eq_of_u_plot = None
        self.create_figure(self.__f)

        # integrate figure in GUI
        self.__plots = FigureCanvasTkAgg(self.__f, master)
        self.__plots.get_tk_widget().pack(side=tk.LEFT)

        # add animation (in order for the plots to be responsive wrt the slider positions)
        self.__anim = FuncAnimation(self.__f, self.refresh_plots, frames=self.__num_steps,
                             interval=1000*self.__delta_t, blit=True, repeat=True)

        # slider for prediction time
        self.create_slider(master, "prediction time [s]", [0, self.__delta_t*self.__num_steps],
                           self.__num_steps, self.adjust_num_steps, self.__delta_t*self.__num_steps)

        # slider for slip angle (state element)
        self.create_slider(master, "slip angle [rad]", [self.__x_bounds[0, 0], self.__x_bounds[0, 1]],
                           self.__x_ticks*10, self.adjust_xi_0, self.__x[0, 0])

        # slider for yaw rate angle (state element)
        self.create_slider(master, "yaw rate [rad/s]", [self.__x_bounds[1, 0], self.__x_bounds[1, 1]],
                           self.__x_ticks*10, self.adjust_xj_0, self.__x[1, 0])

        # slider for steering angle (control input)
        self.create_slider(master, "steering angle [rad]", [self.__u_bounds[0, 0], self.__u_bounds[0, 1]],
                           self.__u_ticks*10, self.adjust_u, self.__u[0, 0])

        # slider for mass (system parameter)
        self.create_slider(master, "mass [kg]", config["m_bounds"], 100, self.adjust_m, self.__m)

        # slider for velocity (system parameter)
        self.create_slider(master, "velocity [m/s]", config["v_bounds"], 100, self.adjust_v, self.__v)

        # slider for moment of inertia (system parameter)
        self.create_slider(master, "moment of inertia [kg m²]", config["theta_bounds"],
                           100, self.adjust_theta, self.__theta)

        # slider for cornering stiffness front (system parameter)
        self.create_slider(master, "cornering stiffness (f) [N/rad]",
                           config["c_f_bounds"], 100, self.adjust_c_f, self.__c_f)

        # slider for cornering stiffness rear (system parameter)
        self.create_slider(master, "cornering stiffness (r) [N/rad]",
                           config["c_r_bounds"], 100, self.adjust_c_r, self.__c_r)

        # slider for front axis to COG (system parameter)
        self.create_slider(master, "front axis to COG [m]",
                           config["l_f_bounds"], 100, self.adjust_l_f, self.__l_f)

        # slider for rear axis to COG (system parameter)
        self.create_slider(master, "rear axis to COG [m]",
                           config["l_r_bounds"], 100, self.adjust_l_r, self.__l_r)

    def recalculate_plot_data(self):
        # init bicycle modle with current parameters
        bicycle_model = BicycleModel(self.__m, self.__v, self.__theta, self.__c_f, self.__c_r, self.__l_f, self.__l_r)

        # equilibrium
        self.__x_eq = bicycle_model.calculate_equilibrium(self.__u)
        self.__x_eq_of_u = bicycle_model.calculate_equilibrium_as_function_of_u(self.__u_bounds, self.__u_ticks)

        # phase space vector field
        self.__xi, self.__xj = np.meshgrid(np.linspace(self.__x_bounds[0, 0], self.__x_bounds[0, 1], self.__x_ticks), np.linspace(
            self.__x_bounds[1, 0], self.__x_bounds[1, 1], self.__x_ticks))
        self.__xi_dot, self.__xj_dot = bicycle_model.generate_vector_field(self.__u, self.__xi, self.__xj)

        # trajectory
        A_d, B_d = bicycle_model.discretize(self.__delta_t)
        self.__time = self.__delta_t * np.arange(0, self.__num_steps+1, 1)
        self.__x_trajectory = bicycle_model.calculate_trajectory(A_d, B_d, self.__x, self.__u, self.__num_steps)

        # laplace of state element step-responses
        # self.__poles = []
        # if self.__x[0] != 0.0 or self.__x[1] != 0.0 or self.__u!= 0.0:
        #     self.__poles = calculate_transfer_function_poles(A).tolist()
        self.__poles = bicycle_model.calculate_transfer_function_poles().tolist()
        if self.__u != 0.0:
            self.__poles.append(0.0+0.0j)

    def create_figure(self, f):
        plt.suptitle("State-Space Representation")
        gs = f.add_gridspec(6, 12)
        ax_phase_space = f.add_subplot(gs[0: 6, 0: 5])
        ax_trajectory_0 = f.add_subplot(gs[0: 3, 5: 9])
        ax_trajectory_1 = f.add_subplot(gs[3: 6, 5: 9])
        ax_laplace_0 = f.add_subplot(gs[0: 3, 9: 12])
        ax_laplace_1 = f.add_subplot(gs[3: 6, 9: 12])

        self.__phase_space_plot = ax_phase_space.quiver(
            self.__xi, self.__xj, self.__xi_dot, self.__xj_dot, angles='xy', scale_units='xy', pivot='mid', zorder=2)
        ax_phase_space.set_xlim((self.__x_bounds[0, 0], self.__x_bounds[0, 1]))
        ax_phase_space.set_ylim((self.__x_bounds[1, 0], self.__x_bounds[1, 1]))
        ax_phase_space.set_xticks(self.__xi[0, :], minor=False)
        ax_phase_space.set_yticks(self.__xj[:, 0], minor=False)
        ax_phase_space.grid(which='major', zorder=1)
        ax_phase_space.set_xlabel('slip angle [rad]')
        ax_phase_space.set_ylabel('yaw rate [rad/s]')

        if self.__x_eq_of_u is not None:
            self.__x_eq_of_u_plot, = ax_phase_space.plot(
                self.__x_eq_of_u[0, :].T, self.__x_eq_of_u[1, :].T, alpha=0.5, color='black')
        else:
            self.__x_eq_of_u_plot, = ax_phase_space.plot([], [], alpha=0.5, color='black')
        self.__trajectory_plot, = ax_phase_space.plot(self.__x_trajectory[0, :].T, self.__x_trajectory[1, :].T, 'b-')
        self.__trajectory_plot_start, = ax_phase_space.plot(self.__x_trajectory[0, 0], self.__x_trajectory[1, 0], 'bo')
        self.__trajectory_plot_end, = ax_phase_space.plot(self.__x_trajectory[0, -1], self.__x_trajectory[1, -1], 'b+')
        if self.__x_eq is not None:
            self.__x_eq_plot, = ax_phase_space.plot(self.__x_eq[0, 0], self.__x_eq[1, 0], 'kx')
        else:
            self.__x_eq_plot, = ax_phase_space.plot([], [], 'kx')

        self.__trajectory_plot_0, = ax_trajectory_0.plot(self.__time, self.__x_trajectory[0, :].T, 'b')
        self.__trajectory_plot_1, = ax_trajectory_1.plot(self.__time, self.__x_trajectory[1, :].T, 'b')
        ax_trajectory_0.grid()
        ax_trajectory_0.set_ylabel('slip angle [rad]')
        ax_trajectory_0.set_xlabel('prediction time [s]')
        ax_trajectory_0.set_ylim((self.__x_bounds[0, 0], self.__x_bounds[0, 1]))
        ax_trajectory_1.grid()
        ax_trajectory_1.set_ylabel('yaw rate [rad/s]')
        ax_trajectory_1.set_xlabel('prediction time [s]')
        ax_trajectory_1.set_ylim((self.__x_bounds[1, 0], self.__x_bounds[1, 1]))

        ax_laplace_0.grid()
        ax_laplace_0.axhline(0, color='black')
        ax_laplace_0.axvline(0, color='black')
        ax_laplace_0.set_xlabel('Re(s)')
        ax_laplace_0.set_ylabel('Im(s) j')
        self.__laplace_0, = ax_laplace_0.plot(np.real(self.__poles), np.imag(self.__poles), 'bx')
        ax_laplace_1.grid()
        ax_laplace_1.axhline(0, color='black')
        ax_laplace_1.axvline(0, color='black')
        ax_laplace_1.set_xlabel('Re(s)')
        ax_laplace_1.set_ylabel('Im(s) j')
        self.__laplace_1, = ax_laplace_1.plot(np.real(self.__poles), np.imag(self.__poles), 'bx')

        f.tight_layout()

    def refresh_plots(self, i):
        # equilibrium
        if self.__x_eq_of_u is not None:
            self.__x_eq_of_u_plot.set_data(self.__x_eq_of_u[0, :].T, self.__x_eq_of_u[1, :].T)
        else:
            self.__x_eq_of_u_plot.set_data([], [])
        if self.__x_eq is not None:
            self.__x_eq_plot.set_data(self.__x_eq[0, 0], self.__x_eq[1, 0])
        else:
            self.__x_eq_plot.set_data([], [])

        # phase space vector field
        self.__phase_space_plot.set_UVC(self.__xi_dot, self.__xj_dot)

        # trajectory
        self.__trajectory_plot.set_data(self.__x_trajectory[0, :].T, self.__x_trajectory[1, :].T)
        self.__trajectory_plot_0.set_data(self.__time, self.__x_trajectory[0, :].T)
        self.__trajectory_plot_1.set_data(self.__time, self.__x_trajectory[1, :].T)
        self.__trajectory_plot_start.set_data(self.__x_trajectory[0, 0], self.__x_trajectory[1, 0])
        self.__trajectory_plot_end.set_data(self.__x_trajectory[0, -1], self.__x_trajectory[1, -1])

        # laplace of state element step-responses
        self.__laplace_0.set_data(np.real(self.__poles), np.imag(self.__poles))
        self.__laplace_1.set_data(np.real(self.__poles), np.imag(self.__poles))

        return (self.__phase_space_plot, self.__trajectory_plot, self.__trajectory_plot_start,
                self.__trajectory_plot_end, self.__trajectory_plot_0, self.__trajectory_plot_1, self.__laplace_0,
                self.__laplace_1, self.__x_eq_plot, self.__x_eq_of_u_plot)

    def create_slider(self, master, sl_name, sl_bounds, sl_steps, sl_cmd, sl_init_val, sl_length=200):
        sl_xj_0 = tk.Scale(master, label=sl_name, from_=sl_bounds[0], to=sl_bounds[1], resolution=(
            sl_bounds[1]-sl_bounds[0])/float(sl_steps), orient=tk.HORIZONTAL, command=sl_cmd, length=sl_length)
        sl_xj_0.set(sl_init_val)
        sl_xj_0.pack()

    def button_method(self):
        # run this when button click to close window
        self.master.destroy()

    def adjust_num_steps(self, val):
        self.__num_steps = int(float(val)/self.__delta_t)
        self.recalculate_plot_data()

    def adjust_xi_0(self, val):
        self.__x[0, 0] = float(val)
        self.recalculate_plot_data()

    def adjust_xj_0(self, val):
        self.__x[1, 0] = float(val)
        self.recalculate_plot_data()

    def adjust_u(self, val):
        self.__u[0, 0] = float(val)
        self.recalculate_plot_data()

    def adjust_m(self, val):
        self.__m = float(val)
        self.recalculate_plot_data()

    def adjust_v(self, val):
        self.__v = float(val)
        self.recalculate_plot_data()

    def adjust_theta(self, val):
        self.__theta = float(val)
        self.recalculate_plot_data()

    def adjust_c_f(self, val):
        self.__c_f = float(val)
        self.recalculate_plot_data()

    def adjust_c_r(self, val):
        self.__c_r = float(val)
        self.recalculate_plot_data()

    def adjust_l_f(self, val):
        self.__l_f = float(val)
        self.recalculate_plot_data()

    def adjust_l_r(self, val):
        self.__l_r = float(val)
        self.recalculate_plot_data()

if __name__ == "__main__":
    # parse arguments
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", required=True, help="directory to configuration yaml-file")
    args = parser.parse_args()

    # Read YAML file
    with open(args.config, 'r') as stream:
        config = yaml.safe_load(stream)

    # launch app
    root = tk.Tk()
    root.configure(background='white')
    app = BicycleModelApp(root, config)
    root = tk.mainloop()


import logging

import requests
from bs4 import BeautifulSoup, Tag
from sqlalchemy.orm import Session

from app import crud, models, const
from app.api import deps

logging.basicConfig(level=logging.WARNING)

def collect_saints_cathedral_saints(cathedral_saints_title: str) -> list[str]:
    r = requests.get(const.AzbykaUrl.GET_SAINT_BY_SLUG + cathedral_saints_title)
    saints_data: Tag = BeautifulSoup(r.text, 'lxml').find(
        lambda tag: tag.name == 'h2'
                    and 'Список Святых' in tag.text
    ).next_sibling.next_sibling
    saints_slugs: list[str] = [saint_slug['href'].replace('/days/sv-', '') for saint_slug in
                               saints_data.find_all('a')]
    return saints_slugs

def main(db: Session):
    # saints_slugs: list[str] = collect_saints_cathedral_saints('sobor-slavnyh-i-vsehvalnyh-12-ti-apostolov')
    saints_slugs: list[str] = collect_saints_cathedral_saints('sobor-70-ti-apostolov')
    # saints_slugs: list[str] = collect_saints_cathedral_saints('sobor-vseh-prepodobnyh-otcov-v-podvige-prosijavshih')
    for saint_slug in saints_slugs:
        saint: models.Saint | None = crud.get_saint(db, saint_slug)
        if not saint:
            print('-' * 10, saint_slug)
        else:
            print(saint.slug)
    print(len(saints_slugs))

if __name__ == '__main__':
    db: Session = next(deps.get_db())
    main(db)


from tkinter import *
from tkinter import ttk
import tkinter as tk

import Util as u
import tkinter.messagebox as msg
import configparser as cp

class DeviceConfig():

	def __init__(self, name, device):
		self.name = name
		self.device = device
		self.ini_file = "conf/"+self.device.dev_sn+".ini"
		#self.ini_file = "conf/global.ini"
		self.ini_elements = {}

	def parse_ini_file(self, ini_file):
		self.active_ini = cp.ConfigParser()
		self.active_ini.read(ini_file)
		self.active_ini_filename = ini_file
		self.display_section_contents()

	def reset_config(self):
		u.copyfile("conf/global.ini", self.ini_file)
		return

	def save_config(self, event=None):
		for key in self.active_ini["SETTING"]:
			self.active_ini["SETTING"][key] = self.ini_elements[key].get()
		with open(self.active_ini_filename, "w") as ini_file:
			self.active_ini.write(ini_file)
		msg.showinfo("Saved", "File Saved Successfully")

	def display_section_contents(self, event=None):

		i = 0
		for key in self.active_ini["SETTING"]:
			new_label = ttk.Label(self.root, text=key, font=(None, 12)).grid(column=0, row=i, sticky=W)
			#print(key+" "+self.active_ini["SETTING"][key])
			value = self.active_ini["SETTING"][key]
			ini_element = ttk.Entry(self.root)
			ini_element.grid(column=1, row=i)
			ini_element.insert(0, value)
			self.ini_elements[key] = ini_element
			i = i+1

		reset_button = tk.Button(self.root, text="Reset Config", command=self.reset_config)
		reset_button.grid(column=0, row=i, sticky=E)

		save_button = tk.Button(self.root, text="Save Config", command=self.save_config)
		save_button.grid(column=1, row=i, sticky=E)
		return

	def initUI(self):
		root = Tk()
		root.title(self.name + " " + self.device.dev_sn)
		root.geometry("300x400+300+200")
		root.resizable(width=False, height=False)
		self.root = root

		if not self.ini_file:
			u.copyfile("conf/global.ini", self.ini_file)
		self.parse_ini_file(self.ini_file)

		root.mainloop()

if __name__ == '__main__':
	obj = DeviceConfig("Config Device")
	obj.initUI()


import sys
sys.dont_write_bytecode = True
import glob
import numpy as np

import utils

DATA_PATH = "/home/satyarth934/data/nasa_impact/hurricanes/*/*"

def main():
    dims = (448, 448, 3)

    # Dataloader creation and test
    img_paths = glob.glob(DATA_PATH)
    print("len(img_paths):", len(img_paths))
    img_paths = utils.getUsableImagePaths(image_paths=img_paths, data_type="png")
    print("---------")
    print("len(usable img_paths):", len(img_paths))
    f = open("hurricane_input.txt", 'w')
    f.writelines("\n".join(img_paths))
    f.close()

if __name__ == "__main__":
    main()


import os
from onsset import *
import pandas as pd
import tkinter as tk
from tkinter import filedialog, messagebox
from openpyxl import load_workbook

root = tk.Tk()
root.withdraw()
root.attributes("-topmost", True)

home_repo_folder = 'D:/OneDrive - FONDAZIONE ENI ENRICO MATTEI/Current papers/Latent demand air cooling/cooling_electricity_SSA/onsset/'
specs_path = home_repo_folder + 'specs_SSA.xlsx'
countries = ['SSA']
SpecsData = pd.read_excel(specs_path, sheet_name='SpecsData')
base_dir = home_repo_folder + 'SSA.csv'
output_dir = home_repo_folder + 'cali_SSA'

print('\n --- Calibrating --- \n')

for country in countries:
	print(country)
	settlements_in_csv = base_dir
	settlements_out_csv = output_dir + '.csv'

	sett = pd.read_csv(base_dir)

	sett['PerCapitaDemand'] = sett['PerHHD']

	sett.to_csv(base_dir)

	onsseter = SettlementProcessor(settlements_in_csv)

	start_year = int(SpecsData.loc[0, SPE_START_YEAR])
	end_year = int(SpecsData.loc[0, SPE_END_YEAR])
	project_life = end_year - start_year

	pop_future_high = int(SpecsData.loc[0, 'PopEndYearHigh'])
	pop_future_low = int(SpecsData.loc[0, 'PopEndYearLow'])
	pop_actual = int(SpecsData.loc[0, SPE_POP])

	urban_growth_high = pop_future_high / pop_actual
	rural_growth_high = pop_future_high / pop_actual

	yearly_urban_growth_rate_high = urban_growth_high ** (1 / project_life)
	yearly_rural_growth_rate_high = rural_growth_high ** (1 / project_life)

	urban_growth_low = pop_future_low / pop_actual
	rural_growth_low = pop_future_low / pop_actual

	yearly_urban_growth_rate_low = urban_growth_low ** (1 / project_life)
	yearly_rural_growth_rate_low = rural_growth_low ** (1 / project_life)

	onsseter.condition_df()
	onsseter.grid_penalties()
	onsseter.calc_wind_cfs()

	pop_actual = SpecsData.loc[0, SPE_POP]
	pop_future_high = SpecsData.loc[0, SPE_POP_FUTURE + 'High']
	pop_future_low = SpecsData.loc[0, SPE_POP_FUTURE + 'Low']
	urban_current = SpecsData.loc[0, SPE_URBAN]
	urban_future = SpecsData.loc[0, SPE_URBAN_FUTURE]
	start_year = int(SpecsData.loc[0, SPE_START_YEAR])
	end_year = int(SpecsData.loc[0, SPE_END_YEAR])

	elec_actual = SpecsData.loc[0, SPE_ELEC]
	elec_actual_urban = SpecsData.loc[0, SPE_ELEC_URBAN]
	elec_actual_rural = SpecsData.loc[0, SPE_ELEC_RURAL]
	pop_tot = SpecsData.loc[0, SPE_POP]

	num_people_per_hh_rural = float(SpecsData.iloc[0][SPE_NUM_PEOPLE_PER_HH_RURAL])
	num_people_per_hh_urban = float(SpecsData.iloc[0][SPE_NUM_PEOPLE_PER_HH_URBAN])

	# In case there are limitations in the way grid expansion is moving in a country, this can be reflected through gridspeed.

	# In this case the parameter is set to a very high value therefore is not taken into account.
	urban_modelled = SpecsData.loc[0, SPE_URBAN_MODELLED]
	elec_modelled = SpecsData.loc[0, SPE_ELEC_MODELLED]
	rural_elec_ratio = SpecsData.loc[0, 'rural_elec_ratio_modelled']
	urban_elec_ratio = SpecsData.loc[0, 'urban_elec_ratio_modelled']

	book = load_workbook(specs_path)
	writer = pd.ExcelWriter(specs_path, engine='openpyxl')
	writer.book = book
	# RUN_PARAM: Here the calibrated "specs" data are copied to a new tab called "SpecsDataCalib". This is what will later on be used to feed the model
	SpecsData.to_excel(writer, sheet_name='SpecsDataCalib', index=False)
	writer.save()
	writer.close()

	logging.info('Calibration finished. Results are transferred to the csv file')
	onsseter.df.to_csv(settlements_out_csv, index=False)


#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from typing import Optional
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
from Environment import StochasticWindyGridworld
from Helper import softmax, argmax

class Agent:
    ''' Base class for agents '''

    def __init__(
        self,
        n_states: int,
        n_actions: int,
        learning_rate: float,
        gamma: float
        ) -> None:
        ''' Initialize the agent '''
        self.n_states = n_states
        self.n_actions = n_actions
        self.learning_rate = learning_rate
        self.gamma = gamma
        self.Q = np.zeros((n_states, n_actions))

    def select_action(
        self,
        s: int,
        policy: str = 'egreedy',
        epsilon: Optional[float] = None,
        temp: Optional[float] = None
        ) -> int:

        if policy == 'egreedy':
            if epsilon is None:
                raise ValueError('Provide an epsilon to use egreedy policy')
            a = self.select_egreedy_action(s, epsilon)
        elif policy == 'softmax':
            if temp is None:
                raise ValueError('Provide a temperature to use softmax policy')
            a = self.select_softmax_action(s, temp)
        else:
            raise ValueError('Unknown policy, please use egreedy or softmax')

        return a

    def select_egreedy_action(
        self,
        s: int,
        epsilon: float
        ) -> int:
        ''' Returns an action according to the epsilon-greedy policy '''
        greedy_a = argmax(self.Q[s])
        if np.random.uniform() < epsilon:
            return np.random.choice([a for a in range(self.n_actions) if a != greedy_a])
        return greedy_a

    def select_softmax_action(
        self,
        s: int,
        temp: float
        ) -> int:
        ''' Returns an action according to the softmax policy '''
        return np.random.choice(self.n_actions, p=softmax(self.Q[s], temp))

    def update(self) -> None:
        raise NotImplementedError

    def run_greedy_episode(self, env: StochasticWindyGridworld) -> float:
        ''' Run a greedy episode and returns the average reward per timestep '''
        s = env.reset()
        done = False
        rewards = []
        for _ in range(100):
            a = argmax(self.Q[s])
            s, r, done = env.step(a)
            rewards.append(r)
            if done:
                break
        return np.mean(rewards)

    def plot_policy(self, title: str, filename: str) -> None:
        ''' Plot the policy based on the current state of self.Q
        Each state is represented by an arrow pointing in the direction of the action with the highest value
        The state is plotted as a 2D representation of the environment with a BW heatmap of the greedy Q-values
        '''
        fig, ax = plt.subplots(1, 1, figsize=(10, 6))

        # prepare the Q-value array
        Q_2D = self.Q.reshape(10, 7, self.n_actions)
        Q_2D = np.transpose(Q_2D, (1, 0, 2))
        Q_2D = np.flip(Q_2D, axis=0)

        # plot the heatmap of the greedy Q-values
        greedy_Q: np.ndarray = np.max(Q_2D, axis=2)
        ax.imshow(greedy_Q, cmap='gray')
        # add colorbar with correct range of values under the heatmap
        cbar = ax.figure.colorbar(ax.images[0], ax=ax, location='right', shrink=0.7)
        cbar.ax.set_ylabel('max($Q(s)$)', rotation=-90, va="bottom", fontsize=14)

        # plot the arrows
        greedy_actions = np.argmax(Q_2D, axis=2, keepdims=True)
        for i in range(7):
            for j in range(10):
                # get the greedy action
                a = greedy_actions[i, j, 0]
                # get the direction of the arrow
                dy, dx = [(-.2, 0), (0, .2), (.2, 0), (0, -.2)][a]
                # plot the arrow
                ax.arrow(j, i, dx, dy, head_width=0.1, head_length=0.1, fc='k', ec='k')

        # add red borders around where wind is blowing
        rect1 = Rectangle((2.5, -0.5), 3, 7, linewidth=3, edgecolor='r', facecolor='none')
        rect2 = Rectangle((5.5, -0.5), 2, 7, linewidth=5, edgecolor='r', facecolor='none')
        rect3 = Rectangle((7.5, -0.5), 1, 7, linewidth=3, edgecolor='r', facecolor='none')
        for rect in [rect1, rect2, rect3]:
            ax.add_patch(rect)

        # add green circle where the goal is
        ax.plot([7], [3], 'go', markersize=30)

        # finish the plot
        ax.set_xticks(range(10))
        ax.set_yticks(range(7))
        ax.set_title(title, fontsize=16)
        fig.tight_layout()
        fig.savefig(filename, dpi=300)


from typing import Optional, Sequence, Tuple

import numpy as np
import torch
import torch.nn as nn

from rl_algo_impls.shared.actor import Actor, PiForward, actor_head
from rl_algo_impls.shared.encoder import Encoder
from rl_algo_impls.shared.policy.actor_critic import OnPolicy, Step, clamp_actions
from rl_algo_impls.shared.policy.actor_critic_network import default_hidden_sizes
from rl_algo_impls.shared.policy.critic import CriticHead
from rl_algo_impls.shared.policy.policy import ACTIVATION
from rl_algo_impls.wrappers.vectorable_wrapper import (
    VecEnv,
    VecEnvObs,
    single_action_space,
    single_observation_space,
)

PI_FILE_NAME = "pi.pt"
V_FILE_NAME = "v.pt"

class VPGActor(Actor):
    def __init__(self, feature_extractor: Encoder, head: Actor) -> None:
        super().__init__()
        self.feature_extractor = feature_extractor
        self.head = head

    def forward(self, obs: torch.Tensor, a: Optional[torch.Tensor] = None) -> PiForward:
        fe = self.feature_extractor(obs)
        return self.head(fe, a)

    def sample_weights(self, batch_size: int = 1) -> None:
        self.head.sample_weights(batch_size=batch_size)

    @property
    def action_shape(self) -> Tuple[int, ...]:
        return self.head.action_shape

class VPGActorCritic(OnPolicy):
    def __init__(
        self,
        env: VecEnv,
        hidden_sizes: Optional[Sequence[int]] = None,
        init_layers_orthogonal: bool = True,
        activation_fn: str = "tanh",
        log_std_init: float = -0.5,
        use_sde: bool = False,
        full_std: bool = True,
        squash_output: bool = False,
        **kwargs,
    ) -> None:
        super().__init__(env, **kwargs)
        activation = ACTIVATION[activation_fn]
        obs_space = single_observation_space(env)
        self.action_space = single_action_space(env)
        self.use_sde = use_sde
        self.squash_output = squash_output

        hidden_sizes = (
            hidden_sizes
            if hidden_sizes is not None
            else default_hidden_sizes(obs_space)
        )

        pi_feature_extractor = Encoder(
            obs_space, activation, init_layers_orthogonal=init_layers_orthogonal
        )
        pi_head = actor_head(
            self.action_space,
            pi_feature_extractor.out_dim,
            tuple(hidden_sizes),
            init_layers_orthogonal,
            activation,
            log_std_init=log_std_init,
            use_sde=use_sde,
            full_std=full_std,
            squash_output=squash_output,
        )
        self.pi = VPGActor(pi_feature_extractor, pi_head)

        v_feature_extractor = Encoder(
            obs_space, activation, init_layers_orthogonal=init_layers_orthogonal
        )
        v_head = CriticHead(
            v_feature_extractor.out_dim,
            tuple(hidden_sizes),
            activation=activation,
            init_layers_orthogonal=init_layers_orthogonal,
        )
        self.v = nn.Sequential(v_feature_extractor, v_head)

    def value(self, obs: VecEnvObs) -> np.ndarray:
        o = self._as_tensor(obs)
        with torch.no_grad():
            v = self.v(o)
        return v.cpu().numpy()

    def step(self, obs: VecEnvObs, action_masks: Optional[np.ndarray] = None) -> Step:
        assert (
            action_masks is None
        ), f"action_masks not currently supported in {self.__class__.__name__}"
        o = self._as_tensor(obs)
        with torch.no_grad():
            pi, _, _ = self.pi(o)
            a = pi.sample()
            logp_a = pi.log_prob(a)

            v = self.v(o)

        a_np = a.cpu().numpy()
        clamped_a_np = clamp_actions(a_np, self.action_space, self.squash_output)
        return Step(a_np, v.cpu().numpy(), logp_a.cpu().numpy(), clamped_a_np)

    def act(
        self,
        obs: np.ndarray,
        deterministic: bool = True,
        action_masks: Optional[np.ndarray] = None,
    ) -> np.ndarray:
        assert (
            action_masks is None
        ), f"action_masks not currently supported in {self.__class__.__name__}"
        if not deterministic:
            return self.step(obs).clamped_a
        else:
            o = self._as_tensor(obs)
            with torch.no_grad():
                pi, _, _ = self.pi(o)
                a = pi.mode
            return clamp_actions(a.cpu().numpy(), self.action_space, self.squash_output)

    def load(self, path: str) -> None:
        super().load(path)
        self.reset_noise()

    def reset_noise(self, batch_size: Optional[int] = None) -> None:
        self.pi.sample_weights(
            batch_size=batch_size if batch_size else self.env.num_envs
        )

    @property
    def action_shape(self) -> Tuple[int, ...]:
        return self.pi.action_shape


# main.py
import argparse
import json
from tabulate import tabulate
import pandas as pd

from core.banner import banner
from core.export import export_results
from core.scanner import run_scan

def main():
    banner()

    parser = argparse.ArgumentParser()
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("-t", "--target", help="Single IP address to scan")
    group.add_argument("-f", "--file", help="Text file containing IP addresses to scan")
    parser.add_argument("-oJ", "--output", help="Output JSON filename")
    parser.add_argument("-mT", "--multithread", type=int, metavar="N", help="Enable multithreading with N threads")
    args = parser.parse_args()

    if args.target:
        ip_list = [args.target]
    elif args.file:
        with open(args.file, "r") as f:
            ip_list = f.readlines()

    if not ip_list:
        print("[INFO] No IP addresses found in the input.")
        return

    try:
        results = run_scan(ip_list, args.multithread)

        if results:
            print_results(results)
            export_results(results, args.output)
        else:
            print("[INFO] No C2 servers detected.")
    except Exception as e:
        print("[ERROR] An error occurred:", e)

def print_results(results):
    table_data = []
    headers = ["STATUS", "IP ADDRESS", "CN", "ASN", "PAGE TITLE", "FULL URL"]

    for result in results:
        table_data.append([
            result.get("STATUS", ""),
            result.get("IP ADDRESS", ""),
            result.get("CN", ""),
            result.get("ASN", ""),
            result.get("PAGE TITLE", ""),
            result.get("FULL URL", "")
        ])

    df = pd.DataFrame(table_data, columns=headers)

    pd.set_option("display.max_colwidth", None)
    pd.set_option("display.width", None)
    pd.set_option("display.max_rows", None)
    pd.set_option("display.max_columns", None)

    if not df.empty:
        df = df.reset_index(drop=True)
        df.index += 1
        print(tabulate(df, headers, tablefmt="simple"))
    else:
        print("[INFO] No C2 servers detected.")

if __name__ == "__main__":
    main()


import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from xgboost import XGBClassifier
import xgboost as xgb
import optuna
from sklearn.model_selection import cross_val_score
from sklearn.model_selection import validation_curve
from sklearn.model_selection import KFold
import time
from sklearn.metrics import accuracy_score

df=pd.read_csv('data/alzheimer_full.csv',index_col='IDENTIFIER')

#df=df.drop(columns='GeneName')
train_set,test_set=train_test_split(df,test_size=0.2,random_state=4)#trainとtestを1:4の割合で分割する、random_stateを固定することでいつも同じ抽出方法

X_train=train_set.drop('Label',axis=1)
y_train=train_set['Label']#Labelをもとに判定したい、訓練セット

X_test=test_set.drop('Label',axis=1)
y_test=test_set['Label']

#X_train=X_train.values
#y=y_train.values

dtrain=xgb.DMatrix(X_train,label=y_train)#xgbが判定できるような特別な配列に変換
dtest=xgb.DMatrix(X_test,label=y_test)

seed=42
model = XGBClassifier(booster='gbtree', objective='binary:logistic',random_state=seed, n_estimators=10000)  # チューニング前のモデル,2値分類なので、objectiveはbinary:logistic

cv = KFold(n_splits=5, shuffle=True, random_state=seed)  # KFoldでクロスバリデーション分割指定

fit_params = {'verbose': 0,  # 学習中のコマンドライン出力
'early_stopping_rounds': 10,  # 学習時、評価指標がこの回数連続で改善しなくなった時点でストップ
'eval_metric': 'rmse',  # early_stopping_roundsの評価指標
'eval_set': [(X_train, y_train)]  # early_stopping_roundsの評価指標算出用データ
}

scoring='neg_log_loss'#2値分類の一般的な評価指標はlogloss

cv_params = {'subsample': [0, 0.1, 0.2, 0.3, 0.4, 0.6, 0.7, 0.8, 0.9, 1.0],
            'colsample_bytree': [0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0],
            'reg_alpha': [0, 0.0001, 0.001, 0.01, 0.03, 0.1, 0.3, 1.0],
            'reg_lambda': [0, 0.0001, 0.001, 0.01, 0.03, 0.1, 0.3, 1.0],
            'learning_rate': [0, 0.0001, 0.001, 0.01, 0.03, 0.1, 0.3, 1.0],
            'min_child_weight': [1, 3, 5, 7, 9, 11, 13, 15],
            'max_depth': [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
            'gamma': [0, 0.0001, 0.001, 0.01, 0.03, 0.1, 0.3, 1.0]
            }
param_scales = {'subsample': 'linear',
                'colsample_bytree': 'linear',
                'reg_alpha': 'log',
                'reg_lambda': 'log',
                'learning_rate': 'log',
                'min_child_weight': 'linear',
                'max_depth': 'linear',
                'gamma': 'log'
                }
# 検証曲線のプロット（パラメータ毎にプロット）
for i, (k, v) in enumerate(cv_params.items()):
    train_scores, valid_scores = validation_curve(estimator=model,
                                                X=X_train, y=y_train,
                                                param_name=k,
                                                param_range=v,
                                                fit_params=fit_params,
                                                cv=cv, scoring=scoring,
                                                n_jobs=64)
    # 学習データに対するスコアの平均±標準偏差を算出
    train_mean = np.mean(train_scores, axis=1)
    train_std  = np.std(train_scores, axis=1)
    train_center = train_mean
    train_high = train_mean + train_std
    train_low = train_mean - train_std
    # テストデータに対するスコアの平均±標準偏差を算出
    valid_mean = np.mean(valid_scores, axis=1)
    valid_std  = np.std(valid_scores, axis=1)
    valid_center = valid_mean
    valid_high = valid_mean + valid_std
    valid_low = valid_mean - valid_std
    # training_scoresをプロット
    fig = plt.figure()
    plt.plot(v, train_center, color='blue', marker='o', markersize=5, label='training score')
    plt.fill_between(v, train_high, train_low, alpha=0.15, color='blue')
    # validation_scoresをプロット
    plt.plot(v, valid_center, color='green', linestyle='--', marker='o', markersize=5, label='validation score')
    plt.fill_between(v, valid_high, valid_low, alpha=0.15, color='green')
    # スケールをparam_scalesに合わせて変更
    plt.xscale(param_scales[k])
    # 軸ラベルおよび凡例の指定
    plt.xlabel(k)  # パラメータ名を横軸ラベルに
    plt.ylabel(scoring)  # スコア名を縦軸ラベルに
    plt.legend(loc='lower right')  # 凡例
    # グラフを描画
    plt.show()
    fig.savefig("graph/alzheimer_full%03.f"%(i)+".png")


import astropy.units as u

import numpy as np
from scipy.signal import savgol_filter
import logging

def find_regions(data, continuum=None, cap_value=None, smooth=False,
                 rel_tol=1e-2, abs_tol=1e-5):
    if cap_value is not None:
        data = np.array(data)
        data[data > cap_value] = cap_value

    if continuum is None:
        continuum = np.zeros(data.shape)

    data = savgol_filter(data, 49, 3) if smooth else data

    # cind = find_nearest(dispersion, center)
    creg = _get_absorption_regions(data, continuum=continuum, rel_tol=rel_tol,
                                   abs_tol=abs_tol)

    return creg

def _get_absorption_regions(data, continuum, rel_tol, abs_tol):
    # This makes the assumption that the continuum has been normalized to 1
    mask = ~np.isclose(data, continuum, rtol=rel_tol, atol=abs_tol)

    creg = contiguous_regions(mask)

    return creg

def contiguous_regions(condition):
    """
    Finds contiguous True regions of the boolean array "condition". Returns
    a 2D array where the first column is the start index of the region and the
    second column is the end index.
    """
    # Find the indices of changes in "condition"
    d = np.diff(condition)
    idx, = d.nonzero()

    # We need to start things after the change in "condition". Therefore,
    # we'll shift the index by 1 to the right.
    # idx += 1

    if condition[0]:
        # If the start of condition is True prepend a 0
        idx = np.r_[0, idx]

    if condition[-1]:
        # If the end of condition is True, append the length of the array
        idx = np.r_[idx, condition.size - 1]

    # Reshape the result into two columns
    idx.shape = (-1, 2)

    return idx


import collections.abc as collections
import numbers
import os
import stat
import tempfile
from abc import ABCMeta, abstractmethod
from collections import OrderedDict
from pathlib import Path
from typing import Any, Callable, cast, Dict, List, Mapping, NamedTuple, Optional, Union

import torch
import torch.nn as nn
from packaging.version import Version

if Version(torch.__version__) >= Version("1.9.0"):
    from torch.distributed.optim import ZeroRedundancyOptimizer

    HAVE_ZERO = True
else:
    HAVE_ZERO = False

import ignite.distributed as idist
from ignite.base import Serializable
from ignite.engine import Engine, Events
from ignite.utils import _tree_apply2, _tree_map

__all__ = ["Checkpoint", "DiskSaver", "ModelCheckpoint", "BaseSaveHandler"]

class BaseSaveHandler(metaclass=ABCMeta):
    """Base class for save handlers

    Methods to override:

    - :meth:`~ignite.handlers.checkpoint.BaseSaveHandler.__call__`
    - :meth:`~ignite.handlers.checkpoint.BaseSaveHandler.remove`

    Note:
        In derived class, please, make sure that in distributed configuration overridden methods are called by a single
        process. Distributed configuration on XLA devices should be treated slightly differently: for saving checkpoint
        with `xm.save() <https://pytorch.org/xla/release/1.5/index.html#torch_xla.core.xla_model.save>`_  all processes
        should pass into the function. Otherwise, application gets stuck.

    """

    @abstractmethod
    def __call__(self, checkpoint: Mapping, filename: str, metadata: Optional[Mapping] = None) -> None:
        """Method to save `checkpoint` with `filename`. Additionally, metadata dictionary is provided.

        Metadata contains:

        - `basename`: file prefix (if provided) with checkpoint name, e.g. `epoch_checkpoint`.
        - `score_name`: score name if provided, e.g `val_acc`.
        - `priority`: checkpoint priority value (higher is better), e.g. `12` or `0.6554435`

        Args:
            checkpoint: checkpoint dictionary to save.
            filename: filename associated with checkpoint.
            metadata: metadata on checkpoint to save.

        """

    @abstractmethod
    def remove(self, filename: str) -> None:
        """Method to remove saved checkpoint.

        Args:
            filename: filename associated with checkpoint.

        """

class Checkpoint(Serializable):
    """Checkpoint handler can be used to periodically save and load objects which have attribute
    ``state_dict/load_state_dict``. This class can use specific save handlers to store on the disk or a cloud
    storage, etc. The Checkpoint handler (if used with :class:`~ignite.handlers.DiskSaver`) also handles automatically
    moving data on TPU to CPU before writing the checkpoint.

    Args:
        to_save: Dictionary with the objects to save. Objects should have implemented ``state_dict`` and
            ``load_state_dict`` methods. If contains objects of type torch `DistributedDataParallel`_ or
            `DataParallel`_, their internal wrapped model is automatically saved (to avoid additional key ``module.`` in
            the state dictionary).
        save_handler: String, function or callable object.
            used to save engine and other provided objects. Function receives two objects: checkpoint as a dictionary
            and filename. If ``save_handler`` is callable class, it can
            inherit of :class:`~ignite.handlers.checkpoint.BaseSaveHandler` and optionally implement ``remove`` method
            to keep a fixed number of saved checkpoints. In case if user needs to save engine's checkpoint on a disk,
            ``save_handler`` can be defined with :class:`~ignite.handlers.DiskSaver` or a string specifying
            directory name can be passed to ``save_handler``.
        filename_prefix: Prefix for the file name to which objects will be saved. See Note for details.
        score_function: If not None, it should be a function taking a single argument,
            :class:`~ignite.engine.engine.Engine` object, and returning a score (`float`). Objects with highest scores
            will be retained.
        score_name: If ``score_function`` not None, it is possible to store its value using
            ``score_name``. If ``score_function`` is None, ``score_name`` can be used alone to define ``score_function``
            as ``Checkpoint.get_default_score_fn(score_name)`` by default.
        n_saved: Number of objects that should be kept on disk. Older files will be removed. If set to
            `None`, all objects are kept.
        global_step_transform: global step transform function to output a desired global step.
            Input of the function is ``(engine, event_name)``. Output of function should be an integer.
            Default is None, global_step based on attached engine. If provided, uses function output as global_step.
            To setup global step from another engine, please use :meth:`~ignite.handlers.global_step_from_engine`.
        filename_pattern: If ``filename_pattern`` is provided, this pattern will be used to render
            checkpoint filenames. If the pattern is not defined, the default pattern would be used. See Note for
            details.
        include_self: Whether to include the `state_dict` of this object in the checkpoint. If `True`, then
            there must not be another object in ``to_save`` with key ``checkpointer``.
        greater_or_equal: if `True`, the latest equally scored model is stored. Otherwise, the first model.
            Default, `False`.
        save_on_rank: Which rank to save the objects on, in the distributed configuration. If ``save_handler`` is
            string or :class:`~pathlib.Path`, this is also used to instantiate a :class:`~ignite.handlers.DiskSaver`.

    .. _DistributedDataParallel: https://pytorch.org/docs/stable/generated/
        torch.nn.parallel.DistributedDataParallel.html
    .. _DataParallel: https://pytorch.org/docs/stable/generated/torch.nn.DataParallel.html

    Note:
        This class stores a single file as a dictionary of provided objects to save.
        The filename is defined by ``filename_pattern`` and by default has the following
        structure: ``{filename_prefix}_{name}_{suffix}.{ext}`` where

        - ``filename_prefix`` is the argument passed to the constructor,
        - `name` is the key in ``to_save`` if a single object is to store, otherwise `name` is "checkpoint".
        - `suffix` is composed as following ``{global_step}_{score_name}={score}``.

    +----------------+------------+-----------------------+----------------------------------------------+
    | score_function | score_name | global_step_transform |  suffix                                      |
    +================+============+=======================+==============================================+
    |      None      |   None     |        None           | ``{engine.state.iteration}``                 |
    +----------------+------------+-----------------------+----------------------------------------------+
    |       X        |   None     |        None           | ``{score}``                                  |
    +----------------+------------+-----------------------+----------------------------------------------+
    |       X        |   None     |         X             | ``{global_step}_{score}``                    |
    +----------------+------------+-----------------------+----------------------------------------------+
    |       X        |    X       |         X             | ``{global_step}_{score_name}={score}``       |
    +----------------+------------+-----------------------+----------------------------------------------+
    |      None      |   None     |         X             | ``{global_step}``                            |
    +----------------+------------+-----------------------+----------------------------------------------+
    |       X        |    X       |        None           | ``{score_name}={score}``                     |
    +----------------+------------+-----------------------+----------------------------------------------+

    Above `global_step` defined by the output of `global_step_transform` and `score` defined by the output
    of `score_function`.

    By default, none of ``score_function``, ``score_name``, ``global_step_transform`` is defined, then suffix is
    setup by attached engine's current iteration. The filename will be
    `{filename_prefix}_{name}_{engine.state.iteration}.{ext}`.

    For example, ``score_name="neg_val_loss"`` and ``score_function`` that returns `-loss` (as objects with highest
    scores will be retained), then saved filename will be ``{filename_prefix}_{name}_neg_val_loss=-0.1234.pt``.

    Note:
        If ``filename_pattern`` is given, it will be used to render the filenames. ``filename_pattern`` is a string
        that can contain ``{filename_prefix}``, ``{name}``, ``{score}``, ``{score_name}`` and ``{global_step}`` as
        templates.

        For example, let ``filename_pattern="{global_step}-{name}-{score}.pt"`` then the saved filename will be
        ``30000-checkpoint-94.pt``

        **Warning:** Please, keep in mind that if filename collide with already used one to saved a checkpoint,
        new checkpoint will replace the older one. This means that filename like ``checkpoint.pt`` will be saved
        every call and will always be overwritten by newer checkpoints.

    Note:
        To get the last stored filename, handler exposes attribute ``last_checkpoint``:

        .. code-block:: python

            handler = Checkpoint(...)
            ...
            print(handler.last_checkpoint)
            > checkpoint_12345.pt

    Note:
        This class is distributed configuration-friendly: it is not required to instantiate the class in rank 0 only
        process. This class supports automatically distributed configuration and if used with
        :class:`~ignite.handlers.DiskSaver`, checkpoint is stored by rank 0 process.

    .. warning::

        When running on XLA devices or using :class:`~torch.distributed.optim.ZeroRedundancyOptimizer`, it
        should be run in all processes, otherwise application can get stuck while saving the checkpoint.

        .. code-block:: python

            # Wrong:
            # if idist.get_rank() == 0:
            #     handler = Checkpoint(...)
            #     trainer.add_event_handler(Events.ITERATION_COMPLETED(every=1000), handler)

            # Correct:
            handler = Checkpoint(...)
            trainer.add_event_handler(Events.ITERATION_COMPLETED(every=1000), handler)

    Examples:
        Attach the handler to make checkpoints during training:

        .. code-block:: python

            from ignite.engine import Engine, Events
            from ignite.handlers import Checkpoint

            trainer = ...
            model = ...
            optimizer = ...
            lr_scheduler = ...

            to_save = {'model': model, 'optimizer': optimizer, 'lr_scheduler': lr_scheduler, 'trainer': trainer}

            if (checkpoint_iters):
                # A: Output is "checkpoint_<iteration>.pt"
                handler = Checkpoint(
                    to_save, '/tmp/models', n_saved=2
                )
                trainer.add_event_handler(Events.ITERATION_COMPLETED(every=1000), handler)
            else:
                # B:Output is "checkpoint_<epoch>.pt"
                gst = lambda *_: trainer.state.epoch
                handler = Checkpoint(
                    to_save, '/tmp/models', n_saved=2, global_step_transform=gst
                )
                trainer.add_event_handler(Events.EPOCH_COMPLETED, handler)

            trainer.run(data_loader, max_epochs=6)
            > A: ["checkpoint_7000.pt", "checkpoint_8000.pt", ]
            > B: ["checkpoint_5.pt", "checkpoint_6.pt", ]

        Attach the handler to an evaluator to save best model during the training
        according to computed validation metric:

        .. code-block:: python

            from ignite.engine import Engine, Events
            from ignite.handlers import Checkpoint, global_step_from_engine

            trainer = ...
            evaluator = ...
            # Setup Accuracy metric computation on evaluator.
            # evaluator.state.metrics contain 'accuracy',
            # which will be used to define ``score_function`` automatically.
            # Run evaluation on epoch completed event
            # ...

            to_save = {'model': model}
            handler = Checkpoint(
                to_save, '/tmp/models',
                n_saved=2, filename_prefix='best',
                score_name="accuracy",
                global_step_transform=global_step_from_engine(trainer)
            )

            evaluator.add_event_handler(Events.COMPLETED, handler)

            trainer.run(data_loader, max_epochs=10)
            > ["best_model_9_accuracy=0.77.pt", "best_model_10_accuracy=0.78.pt", ]

        Customise the ``save_handler``:

        .. code-block:: python

            handler = Checkpoint(
                to_save, save_handler=DiskSaver('/tmp/models', create_dir=True, **kwargs), n_saved=2
            )

    .. versionchanged:: 0.4.3

        - Checkpoint can save model with same filename.
        - Added ``greater_or_equal`` argument.

    .. versionchanged:: 0.4.7

        - `score_name` can be used to define `score_function` automatically without providing `score_function`.
        - `save_handler` automatically saves to disk if path to directory is provided.
        - `save_on_rank` saves objects on this rank in a distributed configuration.
    """

    Item = NamedTuple("Item", [("priority", int), ("filename", str)])
    _state_dict_all_req_keys = ("_saved",)

    def __init__(
        self,
        to_save: Mapping,
        save_handler: Union[str, Path, Callable, BaseSaveHandler],
        filename_prefix: str = "",
        score_function: Optional[Callable] = None,
        score_name: Optional[str] = None,
        n_saved: Union[int, None] = 1,
        global_step_transform: Optional[Callable] = None,
        filename_pattern: Optional[str] = None,
        include_self: bool = False,
        greater_or_equal: bool = False,
        save_on_rank: int = 0,
    ):
        if not isinstance(to_save, collections.Mapping):
            raise TypeError(f"Argument `to_save` should be a dictionary, but given {type(to_save)}")

        self._check_objects(to_save, "state_dict")

        if include_self:
            if not isinstance(to_save, collections.MutableMapping):
                raise TypeError(
                    f"If `include_self` is True, then `to_save` must be mutable, but given {type(to_save)}."
                )

            if "checkpointer" in to_save:
                raise ValueError(f"Cannot have key 'checkpointer' if `include_self` is True: {to_save}")

        if not (
            isinstance(save_handler, str)
            or isinstance(save_handler, Path)
            or callable(save_handler)
            or isinstance(save_handler, BaseSaveHandler)
        ):
            raise TypeError(
                "Argument `save_handler` should be a string or Path object or callable or inherit from BaseSaveHandler"
            )

        if global_step_transform is not None and not callable(global_step_transform):
            raise TypeError(f"global_step_transform should be a function, got {type(global_step_transform)} instead.")

        self.to_save = to_save
        self.filename_prefix = filename_prefix
        if isinstance(save_handler, str) or isinstance(save_handler, Path):
            self.save_handler = DiskSaver(save_handler, create_dir=True, save_on_rank=save_on_rank)
        else:
            self.save_handler = save_handler  # type: ignore
        self.score_function = score_function
        self.score_name = score_name
        if self.score_name is not None and self.score_function is None:
            self.score_function = self.get_default_score_fn(self.score_name)
        self.n_saved = n_saved
        self.ext = "pt"
        self.global_step_transform = global_step_transform
        self.filename_pattern = filename_pattern
        self._saved: List["Checkpoint.Item"] = []
        self.include_self = include_self
        self.greater_or_equal = greater_or_equal
        self.save_on_rank = save_on_rank

    def _get_filename_pattern(self, global_step: Optional[int]) -> str:
        if self.filename_pattern is None:
            filename_pattern = self.setup_filename_pattern(
                with_prefix=len(self.filename_prefix) > 0,
                with_score=self.score_function is not None,
                with_score_name=self.score_name is not None,
                with_global_step=global_step is not None,
            )
        else:
            filename_pattern = self.filename_pattern
        return filename_pattern

    def reset(self) -> None:
        """Method to reset saved checkpoint names.

        Use this method if the engine will independently run multiple times:

        .. code-block:: python

            from ignite.handlers import Checkpoint

            trainer = ...
            checkpointer = Checkpoint(...)

            trainer.add_event_handler(Events.COMPLETED, checkpointer)
            trainer.add_event_handler(Events.STARTED, checkpointer.reset)

            # fold 0
            trainer.run(data0, max_epochs=max_epochs)
            print("Last checkpoint:", checkpointer.last_checkpoint)

            # fold 1
            trainer.run(data1, max_epochs=max_epochs)
            print("Last checkpoint:", checkpointer.last_checkpoint)

        .. versionadded:: 0.4.3
        """
        self._saved = []

    @property
    def last_checkpoint(self) -> Optional[Union[str, Path]]:
        if len(self._saved) < 1:
            return None

        if not isinstance(self.save_handler, DiskSaver):
            return self._saved[-1].filename

        return self.save_handler.dirname / self._saved[-1].filename

    def _check_lt_n_saved(self, or_equal: bool = False) -> bool:
        if self.n_saved is None:
            return True
        return len(self._saved) < self.n_saved + int(or_equal)

    def _compare_fn(self, new: Union[int, float]) -> bool:
        if self.greater_or_equal:
            return new >= self._saved[0].priority
        else:
            return new > self._saved[0].priority

    def __call__(self, engine: Engine) -> None:
        global_step = None
        if self.global_step_transform is not None:
            global_step = self.global_step_transform(engine, engine.last_event_name)

        if self.score_function is not None:
            priority = self.score_function(engine)
            if not isinstance(priority, numbers.Number):
                raise ValueError("Output of score_function should be a number")
        else:
            if global_step is None:
                global_step = engine.state.get_event_attrib_value(Events.ITERATION_COMPLETED)
            priority = global_step

        if self._check_lt_n_saved() or self._compare_fn(priority):
            priority_str = f"{priority}" if isinstance(priority, numbers.Integral) else f"{priority:.4f}"

            checkpoint = self._setup_checkpoint()

            name = "checkpoint"
            if len(checkpoint) == 1:
                for k in checkpoint:
                    name = k
                checkpoint = checkpoint[name]

            filename_pattern = self._get_filename_pattern(global_step)

            filename_dict = {
                "filename_prefix": self.filename_prefix,
                "ext": self.ext,
                "name": name,
                "score_name": self.score_name,
                "score": priority_str if (self.score_function is not None) else None,
                "global_step": global_step,
            }
            filename = filename_pattern.format(**filename_dict)

            metadata = {
                "basename": f"{self.filename_prefix}{'_' * int(len(self.filename_prefix) > 0)}{name}",
                "score_name": self.score_name,
                "priority": priority,
            }

            try:
                index = list(map(lambda it: it.filename == filename, self._saved)).index(True)
                to_remove = True
            except ValueError:
                index = 0
                to_remove = not self._check_lt_n_saved()

            if to_remove:
                item = self._saved.pop(index)
                if isinstance(self.save_handler, BaseSaveHandler):
                    self.save_handler.remove(item.filename)

            self._saved.append(Checkpoint.Item(priority, filename))
            self._saved.sort(key=lambda it: it[0])

            if self.include_self:
                # Now that we've updated _saved, we can add our own state_dict.
                checkpoint["checkpointer"] = self.state_dict()

            try:
                self.save_handler(checkpoint, filename, metadata)
            except TypeError:
                self.save_handler(checkpoint, filename)

    def _setup_checkpoint(self) -> Dict[str, Any]:
        if self.to_save is not None:

            def func(obj: Any, **kwargs: Any) -> Dict:
                if isinstance(obj, (nn.DataParallel, nn.parallel.DistributedDataParallel)):
                    obj = obj.module
                elif HAVE_ZERO and isinstance(obj, ZeroRedundancyOptimizer):
                    obj.consolidate_state_dict(to=self.save_on_rank)
                    if self.save_on_rank != idist.get_rank():
                        return {}
                return obj.state_dict()

            return cast(Dict[str, Any], _tree_map(func, self.to_save))
        return {}

    @staticmethod
    def setup_filename_pattern(
        with_prefix: bool = True, with_score: bool = True, with_score_name: bool = True, with_global_step: bool = True
    ) -> str:
        """Helper method to get the default filename pattern for a checkpoint.

        Args:
            with_prefix: If True, the ``filename_prefix`` is added to the filename pattern:
                ``{filename_prefix}_{name}...``. Default, True.
            with_score: If True, ``score`` is added to the filename pattern: ``..._{score}.{ext}``.
                Default, True. At least one of ``with_score`` and ``with_global_step`` should be True.
            with_score_name: If True, ``score_name`` is added to the filename pattern:
                ``..._{score_name}={score}.{ext}``. If activated, argument ``with_score`` should be
                also True, otherwise an error is raised. Default, True.
            with_global_step: If True, ``{global_step}`` is added to the
                filename pattern: ``...{name}_{global_step}...``.
                At least one of ``with_score`` and ``with_global_step`` should be True.

        Examples:
            .. code-block:: python

                from ignite.handlers import Checkpoint

                filename_pattern = Checkpoint.setup_filename_pattern()

                print(filename_pattern)
                > "{filename_prefix}_{name}_{global_step}_{score_name}={score}.{ext}"

        .. versionadded:: 0.4.3
        """
        filename_pattern = "{name}"

        if not (with_global_step or with_score):
            raise ValueError("At least one of with_score and with_global_step should be True.")

        if with_global_step:
            filename_pattern += "_{global_step}"

        if with_score_name and with_score:
            filename_pattern += "_{score_name}={score}"
        elif with_score:
            filename_pattern += "_{score}"
        elif with_score_name:
            raise ValueError("If with_score_name is True, with_score should be also True")

        if with_prefix:
            filename_pattern = "{filename_prefix}_" + filename_pattern

        filename_pattern += ".{ext}"
        return filename_pattern

    @staticmethod
    def _check_objects(objs: Mapping, attr: str) -> None:
        def func(obj: Any, **kwargs: Any) -> None:
            if not hasattr(obj, attr):
                raise TypeError(f"Object {type(obj)} should have `{attr}` method")

        _tree_map(func, objs)

    @staticmethod
    def load_objects(to_load: Mapping, checkpoint: Union[str, Mapping, Path], **kwargs: Any) -> None:
        """Helper method to apply ``load_state_dict`` on the objects from ``to_load`` using states from ``checkpoint``.

        Args:
            to_load: a dictionary with objects, e.g. `{"model": model, "optimizer": optimizer, ...}`
            checkpoint: a path, a string filepath or a dictionary with state_dicts to load, e.g.
                `{"model": model_state_dict, "optimizer": opt_state_dict}`. If `to_load` contains a single key,
                then checkpoint can contain directly corresponding state_dict.
            kwargs: Keyword arguments accepted for `nn.Module.load_state_dict()`. Passing `strict=False` enables
                the user to load part of the pretrained model (useful for example, in Transfer Learning)

        Examples:
            .. code-block:: python

                import tempfile
                from pathlib import Path

                import torch

                from ignite.engine import Engine, Events
                from ignite.handlers import ModelCheckpoint, Checkpoint

                trainer = Engine(lambda engine, batch: None)

                with tempfile.TemporaryDirectory() as tmpdirname:
                    handler = ModelCheckpoint(tmpdirname, 'myprefix', n_saved=None, create_dir=True)

                    model = torch.nn.Linear(3, 3)
                    optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)

                    to_save = {"weights": model, "optimizer": optimizer}

                    trainer.add_event_handler(Events.EPOCH_COMPLETED(every=2), handler, to_save)
                    trainer.run(torch.randn(10, 1), 5)

                    to_load = to_save
                    checkpoint_fp = Path(tmpdirname) / 'myprefix_checkpoint_40.pt'
                    checkpoint = torch.load(checkpoint_fp)
                    Checkpoint.load_objects(to_load=to_load, checkpoint=checkpoint)

                    # or using a string for checkpoint filepath

                    to_load = to_save
                    checkpoint_fp = Path(tmpdirname) / 'myprefix_checkpoint_40.pt'
                    Checkpoint.load_objects(to_load=to_load, checkpoint=checkpoint_fp)

        Note:
            If ``to_load`` contains objects of type torch `DistributedDataParallel`_ or
            `DataParallel`_, method ``load_state_dict`` will applied to their internal wrapped model (``obj.module``).

        .. _DistributedDataParallel: https://pytorch.org/docs/stable/generated/
            torch.nn.parallel.DistributedDataParallel.html
        .. _DataParallel: https://pytorch.org/docs/stable/generated/torch.nn.DataParallel.html
        """
        if not isinstance(checkpoint, (collections.Mapping, str, Path)):
            raise TypeError(f"Argument checkpoint should be a string or a dictionary, but given {type(checkpoint)}")

        Checkpoint._check_objects(to_load, "load_state_dict")

        if isinstance(checkpoint, (str, Path)):
            checkpoint_obj = torch.load(checkpoint)
        else:
            checkpoint_obj = checkpoint

        def _load_object(obj: Any, chkpt_obj: Any) -> None:
            if isinstance(obj, (nn.DataParallel, nn.parallel.DistributedDataParallel)):
                obj = obj.module

            if isinstance(obj, torch.nn.Module):
                obj.load_state_dict(chkpt_obj, **kwargs)
            else:
                obj.load_state_dict(chkpt_obj)

        if len(to_load) == 1:
            # single object and checkpoint is directly a state_dict
            key, obj = list(to_load.items())[0]
            if key not in checkpoint_obj:
                _load_object(obj, checkpoint_obj)
                return

        _tree_apply2(_load_object, to_load, checkpoint_obj)

    def reload_objects(self, to_load: Mapping, load_kwargs: Optional[Dict] = None, **filename_components: Any) -> None:
        """Helper method to apply ``load_state_dict`` on the objects from ``to_load``. Filename components such as
        name, score and global state can be configured.

        Args:
            to_load: a dictionary with objects, e.g. `{"model": model, "optimizer": optimizer, ...}`
            load_kwargs: Keyword arguments accepted for `nn.Module.load_state_dict()`. Passing `strict=False` enables
                the user to load part of the pretrained model (useful for example, in Transfer Learning)
            filename_components: Filename components used to define the checkpoint file path.
                Keyword arguments accepted are `name`, `score` and `global_state`.

        Examples:
            .. code-block:: python

                import tempfile

                import torch

                from ignite.engine import Engine, Events
                from ignite.handlers import ModelCheckpoint

                trainer = Engine(lambda engine, batch: None)

                with tempfile.TemporaryDirectory() as tmpdirname:
                    checkpoint = ModelCheckpoint(tmpdirname, 'myprefix', n_saved=None, create_dir=True)

                    model = torch.nn.Linear(3, 3)
                    optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)

                    to_save = {"weights": model, "optimizer": optimizer}

                    trainer.add_event_handler(Events.EPOCH_COMPLETED(every=2), checkpoint, to_save)
                    trainer.run(torch.randn(10, 1), 5)

                    to_load = to_save
                    # load checkpoint myprefix_checkpoint_40.pt
                    checkpoint.reload_objects(to_load=to_load, global_step=40)

        Note:
            If ``to_load`` contains objects of type torch `DistributedDataParallel`_ or
            `DataParallel`_, method ``load_state_dict`` will applied to their internal wrapped model (``obj.module``).

        Note:
            This method works only when the ``save_handler`` is of types string,
            :class:`~pathlib.Path` or :class:`~ignite.handlers.checkpoint.DiskSaver`.

        .. _DistributedDataParallel: https://pytorch.org/docs/stable/generated/
            torch.nn.parallel.DistributedDataParallel.html
        .. _DataParallel: https://pytorch.org/docs/stable/generated/torch.nn.DataParallel.html
        """
        if not isinstance(self.save_handler, DiskSaver):
            raise AttributeError(
                f"Checkpoint's `save_handler` should be of type `DiskSaver`, given {type(self.save_handler)}"
            )

        global_step = filename_components.get("global_step", None)

        filename_pattern = self._get_filename_pattern(global_step)

        checkpoint = self._setup_checkpoint()
        name = "checkpoint"
        if len(checkpoint) == 1:
            for k in checkpoint:
                name = k
        name = filename_components.get("name", name)
        score = filename_components.get("score", None)

        filename_dict = {
            "filename_prefix": self.filename_prefix,
            "ext": self.ext,
            "name": name,
            "score_name": self.score_name,
            "score": score,
            "global_step": global_step,
        }

        checkpoint_fp = filename_pattern.format(**filename_dict)

        path = self.save_handler.dirname / checkpoint_fp

        load_kwargs = {} if load_kwargs is None else load_kwargs

        Checkpoint.load_objects(to_load=to_load, checkpoint=path, **load_kwargs)

    def state_dict(self) -> OrderedDict:
        """Method returns state dict with saved items: list of ``(priority, filename)`` pairs.
        Can be used to save internal state of the class.
        """
        # TODO: this method should use _state_dict_all_req_keys
        return OrderedDict([("_saved", [(p, f) for p, f in self._saved])])

    def load_state_dict(self, state_dict: Mapping) -> None:
        """Method replaces internal state of the class with provided state dict data.

        Args:
            state_dict: a dict with "saved" key and list of ``(priority, filename)`` pairs as values.
        """
        super().load_state_dict(state_dict)
        self._saved = [Checkpoint.Item(p, f) for p, f in state_dict["_saved"]]

    @staticmethod
    def get_default_score_fn(metric_name: str, score_sign: float = 1.0) -> Callable:
        """Helper method to get default score function based on the metric name.

        Args:
            metric_name: metric name to get the value from ``engine.state.metrics``.
                Engine is the one to which :class:`~ignite.handlers.checkpoint.Checkpoint` handler is added.
            score_sign: sign of the score: 1.0 or -1.0. For error-like metrics, e.g. smaller is better,
                a negative score sign should be used (objects with larger score are retained). Default, 1.0.

        Examples:
            .. code-block:: python

                from ignite.handlers import Checkpoint

                best_acc_score = Checkpoint.get_default_score_fn("accuracy")

                best_model_handler = Checkpoint(
                    to_save, save_handler, score_name="val_accuracy", score_function=best_acc_score
                )
                evaluator.add_event_handler(Events.COMPLETED, best_model_handler)

            Usage with error-like metric:

            .. code-block:: python

                from ignite.handlers import Checkpoint

                neg_loss_score = Checkpoint.get_default_score_fn("loss", -1.0)

                best_model_handler = Checkpoint(
                    to_save, save_handler, score_name="val_neg_loss", score_function=neg_loss_score
                )
                evaluator.add_event_handler(Events.COMPLETED, best_model_handler)

        .. versionadded:: 0.4.3
        """
        if score_sign not in (1.0, -1.0):
            raise ValueError("Argument score_sign should be 1 or -1")

        def wrapper(engine: Engine) -> float:
            return score_sign * engine.state.metrics[metric_name]

        return wrapper

class DiskSaver(BaseSaveHandler):
    """Handler that saves input checkpoint on a disk.

    Args:
        dirname: Directory path where the checkpoint will be saved
        atomic: if True, checkpoint is serialized to a temporary file, and then
            moved to final destination, so that files are guaranteed to not be damaged
            (for example if exception occurs during saving).
        create_dir: if True, will create directory ``dirname`` if it doesnt exist.
        require_empty: If True, will raise exception if there are any files in the
            directory ``dirname``.
        save_on_rank: The rank on which the checkpoint will be saved. Used in distributed
            configuration.
        kwargs: Accepted keyword arguments for `torch.save` or `xm.save`.

    .. versionchanged:: 0.4.2
        Accept ``kwargs`` for `torch.save` or `xm.save`.

    .. versionchanged:: 0.4.10
        Argument ``save_on_rank`` was added to specify the rank on which checkpoint should be saved.
    """

    def __init__(
        self,
        dirname: Union[str, Path],
        atomic: bool = True,
        create_dir: bool = True,
        require_empty: bool = True,
        save_on_rank: int = 0,
        **kwargs: Any,
    ):
        self.dirname = Path(dirname).expanduser()
        self._atomic = atomic
        self.save_on_rank = save_on_rank

        if idist.get_rank() == save_on_rank:
            self._check_and_setup(self.dirname, create_dir, require_empty)
        self.kwargs = kwargs

    @staticmethod
    def _check_and_setup(dirname: Path, create_dir: bool, require_empty: bool) -> None:
        if create_dir:
            if not dirname.exists():
                dirname.mkdir(parents=True)
        # Ensure that dirname exists
        if not dirname.exists():
            raise ValueError(f"Directory path '{dirname}' is not found")

        if require_empty:
            matched = [fname for fname in os.listdir(dirname) if fname.endswith(".pt")]
            if len(matched) > 0:
                raise ValueError(
                    f"Files {matched} with extension '.pt' are already present "
                    f"in the directory {dirname}. If you want to use this "
                    "directory anyway, pass `require_empty=False`."
                    ""
                )

    def __call__(self, checkpoint: Mapping, filename: str, metadata: Optional[Mapping] = None) -> None:
        path = self.dirname / filename

        if idist.has_xla_support:
            import torch_xla.core.xla_model as xm

            # all tpu procs should enter here as internally performs sync across device
            self._save_func(checkpoint, path, xm.save)
        elif self.save_on_rank == idist.get_rank():
            self._save_func(checkpoint, path, torch.save)

    def _save_func(self, checkpoint: Mapping, path: Path, func: Callable) -> None:
        if not self._atomic:
            func(checkpoint, path, **self.kwargs)
        else:
            tmp = tempfile.NamedTemporaryFile(delete=False, dir=self.dirname)
            tmp_file = tmp.file
            tmp_name = tmp.name
            try:
                func(checkpoint, tmp_file, **self.kwargs)
            except BaseException:
                tmp.close()
                os.remove(tmp_name)
                raise
            else:
                tmp.close()
                os.replace(tmp.name, path)
                # append group/others read mode
                os.chmod(path, os.stat(path).st_mode | stat.S_IRGRP | stat.S_IROTH)

    def remove(self, filename: str) -> None:
        if idist.get_rank() == self.save_on_rank:
            path = self.dirname / filename
            path.unlink()

class ModelCheckpoint(Checkpoint):
    """ModelCheckpoint handler, inherits from :class:`~ignite.handlers.checkpoint.Checkpoint`, can be used
    to periodically save objects to disk only. If needed to store checkpoints to
    another storage type, please consider :class:`~ignite.handlers.checkpoint.Checkpoint`.
    It also provides `last_checkpoint` attribute to show the last saved checkpoint.

    This handler expects two arguments:

        - an :class:`~ignite.engine.engine.Engine` object
        - a `dict` mapping names (`str`) to objects that should be saved to disk.

    See Examples for further details.

    .. warning::

        Behaviour of this class has been changed since v0.3.0.

        There is no more internal counter that has been used to indicate the number of save actions. User could
        see its value `step_number` in the filename, e.g. `{filename_prefix}_{name}_{step_number}.pt`. Actually,
        `step_number` is replaced by current engine's epoch if `score_function` is specified and current iteration
        otherwise.

        A single `pt` file is created instead of multiple files.

    Args:
        dirname: Directory path where objects will be saved.
        filename_prefix: Prefix for the file names to which objects will be saved. See Notes of
            :class:`~ignite.handlers.checkpoint.Checkpoint` for more details.
        score_function: if not None, it should be a function taking a single argument, an
            :class:`~ignite.engine.engine.Engine` object, and return a score (`float`). Objects with highest scores
            will be retained.
        score_name: if ``score_function`` not None, it is possible to store its value using
            `score_name`. See Examples of :class:`~ignite.handlers.checkpoint.Checkpoint` for more details.
        n_saved: Number of objects that should be kept on disk. Older files will be removed. If set to
            `None`, all objects are kept.
        atomic: If True, objects are serialized to a temporary file, and then moved to final
            destination, so that files are guaranteed to not be damaged (for example if exception
            occurs during saving).
        require_empty: If True, will raise exception if there are any files starting with
            ``filename_prefix`` in the directory ``dirname``.
        create_dir: If True, will create directory ``dirname`` if it does not exist.
        global_step_transform: global step transform function to output a desired global step.
            Input of the function is `(engine, event_name)`. Output of function should be an integer.
            Default is None, global_step based on attached engine. If provided, uses function output as global_step.
            To setup global step from another engine, please use :meth:`~ignite.handlers.global_step_from_engine`.
        filename_pattern: If ``filename_pattern`` is provided, this pattern will be used to render
            checkpoint filenames. If the pattern is not defined, the default pattern would be used.
            See :class:`~ignite.handlers.checkpoint.Checkpoint` for details.
        include_self: Whether to include the `state_dict` of this object in the checkpoint. If `True`, then
            there must not be another object in ``to_save`` with key ``checkpointer``.
        greater_or_equal: if `True`, the latest equally scored model is stored. Otherwise, the first model.
            Default, `False`.
        save_on_rank: Which rank to save the objects on, in the distributed configuration. Used to
            instantiate a :class:`~ignite.handlers.DiskSaver` and is also passed to the parent class.
        kwargs: Accepted keyword arguments for `torch.save` or `xm.save` in `DiskSaver`.

    .. versionchanged:: 0.4.2
        Accept ``kwargs`` for `torch.save` or `xm.save`

    .. versionchanged:: 0.4.9
        Accept ``filename_pattern`` and ``greater_or_equal`` for parity
        with :class:`~ignite.handlers.checkpoint.Checkpoint`

    .. versionchanged:: 0.4.10
        Added `save_on_rank` arg to save objects on this rank in a distributed configuration

    Examples:
        .. testcode:: python

            import os
            from ignite.engine import Engine, Events
            from ignite.handlers import ModelCheckpoint
            from torch import nn
            trainer = Engine(lambda engine, batch: None)
            handler = ModelCheckpoint('/tmp/models', 'myprefix', n_saved=2, create_dir=True, require_empty=False)
            model = nn.Linear(3, 3)
            trainer.add_event_handler(Events.EPOCH_COMPLETED(every=2), handler, {'mymodel': model})
            trainer.run([0, 1, 2, 3, 4], max_epochs=6)
            print(sorted(os.listdir('/tmp/models')))
            print(handler.last_checkpoint)

        .. testoutput:: python

            ['myprefix_mymodel_20.pt', 'myprefix_mymodel_30.pt']
            /tmp/models/myprefix_mymodel_30.pt
    """

    def __init__(
        self,
        dirname: Union[str, Path],
        filename_prefix: str = "",
        score_function: Optional[Callable] = None,
        score_name: Optional[str] = None,
        n_saved: Union[int, None] = 1,
        atomic: bool = True,
        require_empty: bool = True,
        create_dir: bool = True,
        global_step_transform: Optional[Callable] = None,
        filename_pattern: Optional[str] = None,
        include_self: bool = False,
        greater_or_equal: bool = False,
        save_on_rank: int = 0,
        **kwargs: Any,
    ):
        disk_saver = DiskSaver(
            dirname,
            atomic=atomic,
            create_dir=create_dir,
            require_empty=require_empty,
            save_on_rank=save_on_rank,
            **kwargs,
        )

        super(ModelCheckpoint, self).__init__(
            to_save={},
            save_handler=disk_saver,
            filename_prefix=filename_prefix,
            score_function=score_function,
            score_name=score_name,
            n_saved=n_saved,
            global_step_transform=global_step_transform,
            filename_pattern=filename_pattern,
            include_self=include_self,
            greater_or_equal=greater_or_equal,
            save_on_rank=save_on_rank,
        )

    @property
    def last_checkpoint(self) -> Optional[Union[str, Path]]:
        if len(self._saved) < 1:
            return None

        if not isinstance(self.save_handler, DiskSaver):
            raise RuntimeError(f"Internal error, save_handler should be DiskSaver, but has {type(self.save_handler)}.")

        return self.save_handler.dirname / self._saved[-1].filename

    def __call__(self, engine: Engine, to_save: Mapping):  # type: ignore
        if len(to_save) == 0:
            raise RuntimeError("No objects to checkpoint found.")

        self._check_objects(to_save, "state_dict")
        self.to_save = to_save
        super(ModelCheckpoint, self).__call__(engine)