# Helper functions for "Bernoulli Line and Bloch Sphere" laboratory
#
# These functions are implement to prevent Pythonic distractions
# from the key tasks of the laboratory. If learners have sufficent
# background the code below is written to be instructional and provides
# links to other sources.
#
#
# James Whitfield (2022)
# Dartmouth College
# Amazon Visiting Academic
#
# AJ Cressman (2022)
# Dartmouth College

import matplotlib.pyplot as plt
import numpy as np
from braket.aws import AwsDevice
from braket.circuits import Circuit, Noise
from braket.devices import LocalSimulator

def run_circuit_task1(circuit,n_shots=1000, device=None, verbose=False, plot=True):
    """run_circuit_task1(circuit,n_shots=1000, device=None, verbose=False, plot=True)
    Runs a given single qubit experiment.

        Args:
            circuit

        Optional args:
            n_shots     Integer number of times to run the circuit. Default: 1000.
            device      The default device is the LocalSimulator
            verbose     Boolean variable for printing extra output. Default: False.
            plot        Boolean to create plot. Default: True.

        Returns:
            counts:     A dictionary with keys "0" and "1" containing the counts.

        Examples:
            >>> circuit = Circuit().h(0)
            >>> run_circuit_task1(circuit)
            >>> run_circuit_task1(circuit,n_shots=100)
            >>> run_circuit_task1(circuit,n_shots=300,verbose=True,plot=False)
    """

    if circuit.qubit_count > 1:
        print("Only single qubit circuits please. Returning zero.")
        return 0

    # We will be using the LocalSimulator (rather than an actual quantum device, or one of AWS's managed simulators).
    # The default device is the "Local State Vector Simulator".
    # For more documentation, see: https://docs.aws.amazon.com/braket/latest/developerguide/braket-devices.html
    if device is None:
        device = LocalSimulator()

    # execute task on simulator
    # Note: this is a very common workflow -- set up circuit, run, get result as measurement counts, plot
    counts = device.run(circuit, shots=n_shots).result().measurement_counts

    # counts is a python dictionary
    # (technically, a Counter: https://docs.python.org/3/library/collections.html#collections.Counter):
    if verbose:
        print(counts)

    # plot measurement counts as percentages
    # plt.bar(x-values, y-values, ...options)
    # More documentation: https://matplotlib.org/stable/api/_as_gen/matplotlib.pyplot.bar.html

    if plot:
        plt.bar(("0", "1"), (100 * counts["0"] / n_shots, 100 * counts["1"] / n_shots))
        plt.ylim([0, 100])
        plt.xlabel("Outcome")
        plt.ylabel("Percent")

    if verbose:
        print("Zero counts:", counts["0"])
        print("One  counts:", counts["1"])

    return counts

def run_bernoulli_task(p,n_shots=None,verbose=False,plot=True):
    """run_bernoulli_task(p,n_shots=None,verbose=False,plot=True)

    Runs a given Bernoulli experiment.

        Args:
            p           Probability of Outcome 0

        Optional args:
            n_shots     Integer number of times to run the circuit. Default: 1000.
            verbose     Boolean to print extra. Default: False.
            plot        Boolean to create plot. Default: True.

        Returns:
            counts:     A dictionary with keys "0" and "1" containing the counts.

        Examples:
            >>> p=0.8
            >>> n=1000
            >>> run_bernoulli_task(p,n)
            >>> run_bernoulli_task(p,n,verbose=True,plot=False)
    """
    if n_shots is None:
        n_shots = 1000

    circuit = Circuit()

    # Define a bit flip noise channel
    # Since the channel is considered noise, it can only cause error with prob 0 to 0.5.
    # Otherwise, it is performing a faulty X gate
    if p > 0.5:
        # apply Z gate to first qubit (counting starts at 0)
        # Note: If we left this out, we would just have an empty circuit, and our simulation would fail!
        circuit.z(0)
        noise = Noise.BitFlip(probability = 1 - p)
    else:
        # apply X gate to first qubit (counting starts at 0)
        circuit.x(0)
        noise = Noise.BitFlip(probability = p)

    circuit.apply_gate_noise(noise)

    # Our usual workflow, with a different device
    # Here we are using the "Local Density Matrix Simulator", because we are running noise instructions.
    # Students: change this back to the default (LocalSimulator()) and see what error message you get.
    device = LocalSimulator("braket_dm")
    counts = device.run(circuit, shots=n_shots).result().measurement_counts

    if(plot):
        plt.bar(("0", "1"), (100 * counts["0"] / n_shots, 100 * counts["1"] / n_shots))
        plt.ylim([0, 100])
        plt.xlabel("outcome")
        plt.ylabel("percent")

    if(verbose):
        print("Zero counts:", counts["0"])
        print("One  counts:", counts["1"])

    return counts

def get_bernoulli_val_to_estimate(seedval=None):
    """get_bernoulli_val_to_estimate(seedval=None)
    This function returns a random Bernoulli parameter p

    **Do not print the value if you are performing an estimation procedure**

    Optional args:
        seedval         A seed value for the random number generator.
                        This ensures that the same random number is generated
                        each time the function is run.
    Return:
        p_unknown       Bernoulli parameter to be estimated
    """

    # Don't dig into this too much...
    # Try to figure it out without printing p!

    if(seedval):
        np.random.seed(seedval)

    p_unknown = np.random.rand()

    return p_unknown

# This function will generate a circuit for a random state.
# Students: write some code to estimate the state of the qubit.
# Be as quantitative as possible.  And a plot would be nice!
def get_target_bias_parameters(seedval=1):
    """get_target_bias_parameters(seedval=1)
    This function returns random vector of unit length

    **Do not print the values if you are performing an estimation procedure**

    Optional args:
        seedval         A seed value for the random number generator.
                        This ensures that the same random number is generated
                        each time the function is run.
    Return:
        x,y,z           Parameters to be estimated
    """
    # Don't dig into this too much...
    # Try to figure it out without printing the circuit parameters!
    np.random.seed(seedval)

    # initial
    xi = np.random.randn()
    yi = np.random.randn()
    zi = np.random.randn()

    R = np.sqrt(xi * xi + yi * yi + zi * zi)

    x = xi / R
    y = yi / R
    z = zi / R

    return x, y, z

def get_circuit_of_quantum_state_to_estimate(seedval=None):
    """get_circuit_of_quantum_state_to_estimate(seedval=None)

    This function returns a circuit to create a random qubit state

    **Do not print the circuit if you are performing an estimation procedure**

    Optional args:
        seedval         A seed value for the random number generator.
                        This ensures that the same random numbers are generated
                        each time the function is run.
    Return:
        cir             A state preparation circuit
    """

    if(seedval):
        np.random.seed(seedval)

    state_prep_circuit = Circuit()
    state_prep_circuit.rx(0, np.random.rand() * 2 * np.pi)
    state_prep_circuit.ry(0, np.random.rand() * 2 * np.pi)

    return state_prep_circuit

def plot_qubit_state(qubit_biases,computed_uncertainties=None):
    """plot_qubit_state(qubit_biases,computed_uncertainties=None)

    Plots a pure state qubit along with uncertainties

    Args:
        qubit_biases                Dictionaries with keys "x", "y", "z".
                                    Values should be normalized to unit vector length.
        computed_uncertainties      Dictionaries with keys "x", "y", "z" giving uncertainies.

    Examples:
        >>> qubit_biases           = {"x": 0.429, "y": -0.01, "z": 0.248}
        >>> computed_uncertainties = {"x": 0.5,   "y": 0.25,  "z": 0.25 }
        >>> plot_qubit_state(qubit_biases)
        >>> plot_qubit_state(qubit_biases,computed_uncertainties)

    """

    # More fun with matplotlib: https://matplotlib.org/3.5.0/gallery/
    fig = plt.figure(dpi=256)
    ax = fig.add_subplot(projection="3d")
    # here we create a 3D mesh of a sphere
    # Students: check to see that this makes sense
    u, v = np.mgrid[0 : 2 * np.pi : 20j, 0 : np.pi : 10j]
    x = np.cos(u) * np.sin(v)
    y = np.sin(u) * np.sin(v)
    z = np.cos(v)
    ax.plot_wireframe(x, y, z, color="red", linewidth=0.1)
    ax.set_title("Vectors+Sphere")
    ax.view_init(45, 35)  # the viewing angle

    # Plot axes
    # We use the `quiver` function to plot arrows
    # More documentation: https://matplotlib.org/3.5.0/api/_as_gen/matplotlib.axes.Axes.quiver.html
    ax.quiver(
        # starting x-position of each arrow
        [0, 0, 0],
        # starting y-position of each arrow
        [0, 0, 0],
        # starting z-position of each arrow
        [0, 0, 0],
        # length in x-direction of each arrow
        [1, 0, 0],
        # length in y-direction of each arrow
        [0, 1, 0],
        # length in z-direction of each arrow
        [0, 0, 1],
        linewidth=1,
    )

    # Plot state
    ax.quiver(
        # starting x-position of each arrow
        [0],
        # starting y-position of each arrow
        [0],
        # starting z-position of each arrow
        [0],
        # length in x-direction of each arrow
        [qubit_biases["x"]],
        # length in y-direction of each arrow
        [qubit_biases["y"]],
        # final z-pos
        [qubit_biases["z"]],
        linewidth=3,
        color="green",
    )

    # Plot uncertainties
    # Note: separated for readability, but can be recombined into one big plot easily
    # start at center of bias and add 0.5 * uncertainty
    ax.quiver(
        # starting x-position of each arrow
        [qubit_biases["x"], 0, 0],
        # starting y-position of each arrow
        [0, qubit_biases["y"], 0],
        # starting z-position of each arrow
        [0, 0, qubit_biases["z"]],
        # length in x-direction of each arrow
        [0.5 * computed_uncertainties["x"], 0, 0],
        # length in y-direction of each arrow
        [0, 0.5 * computed_uncertainties["y"], 0],
        # length in z-direction of each arrow
        [0, 0, 0.5 * computed_uncertainties["z"]],
        color="black",
        linewidth=2,
    )
    # start at center of bias and subtract 0.5 * uncertainty
    ax.quiver(
        # starting x-position of each arrow
        [qubit_biases["x"], 0, 0],
        # starting y-position of each arrow
        [0, qubit_biases["y"], 0],
        # starting z-position of each arrow
        [0, 0, qubit_biases["z"]],
        # length in x-direction of each arrow
        [-0.5 * computed_uncertainties["x"], 0, 0],
        # length in y-direction of each arrow
        [0, -0.5 * computed_uncertainties["y"], 0],
        # length in z-direction of each arrow
        [0, 0, -0.5 * computed_uncertainties["z"]],
        color="black",
        linewidth=2,
    )


import tensorflow as tf
import model as M
import numpy as np

# ------build main model and decoder -----------

def build_model(inph):
	with tf.variable_scope('CapsModel'):
		mod = M.Model(inph,[None,28*28])
		mod.reshape([-1,28,28,1])
		mod.convLayer(11,256,activation=M.PARAM_RELU,pad='VALID')
		mod.primaryCaps(9,8,32,stride=2,activation=M.PARAM_RELU,pad='VALID')
		mod.capsLayer(10,16,3,BSIZE)
		mod.capsDown()
	return mod.get_current_layer()

def build_decoder(inp,labh):
	with tf.variable_scope('Decoder'):
		mod = M.Model(inp,[None,10,16])
		mod.capsMask(labh)
		mod.fcLayer(512,activation=M.PARAM_RELU)
		mod.fcLayer(1024,activation=M.PARAM_RELU)
		mod.fcLayer(28*28,activation=M.PARAM_SIGMOID)
	return mod.get_current_layer()

# -------define margin loss and reconstruction loss ---------
def build_graph():
	inp_holder = tf.placeholder(tf.float32,[None,28*28])
	label_holder = tf.placeholder(tf.float32,[None,10])

	out = build_model(inp_holder)
	decoded = build_decoder(out,label_holder)

	with tf.variable_scope('Length'):
		length = tf.sqrt(tf.reduce_sum(tf.square(out),-1))

	with tf.variable_scope('Loss'):
		marg_loss = label_holder * tf.square(tf.maximum(0.,0.9-length)) + 0.5 * (1-label_holder) * tf.square(tf.maximum(0.,length-0.1))
		marg_loss = tf.reduce_mean(tf.reduce_sum(marg_loss,1),0)
		recon_loss = tf.reduce_mean(tf.square(inp_holder-decoded))

	with tf.variable_scope('Opti'):
		train_step = tf.train.AdamOptimizer(0.0001).minimize(marg_loss+0.0005*784*recon_loss) # multiplier from MSE to SE

	with tf.variable_scope('Accuracy'):
		accuracy = M.accuracy(length,tf.argmax(label_holder,1))

	return inp_holder,label_holder,recon_loss,marg_loss,train_step,accuracy

# ------training-----
from tensorflow.examples.tutorials.mnist import input_data
mnist = input_data.read_data_sets('MNIST_data', one_hot=True)
BSIZE = 100
MAX_ITER = 20000

inph,labh,recon_loss,marg_loss,train_step,accuracy = build_graph()
tensors = [train_step,recon_loss,marg_loss,accuracy]

with tf.Session() as sess:
	tf.summary.FileWriter('./logs/',sess.graph)
	saver = tf.train.Saver()
	M.loadSess('./model/',sess,init=True)
	for iteration in range(MAX_ITER):
		x_train, y_train = mnist.train.next_batch(BSIZE)
		_,lossr,lossm,ac = sess.run([train_step,recon_loss,marg_loss,accuracy], feed_dict={inph:x_train, labh:y_train})
		if iteration % 10 ==0:
			print('Iter:',iteration,'\tLossr:',lossr,'\tLossm:',lossm,'\tAcc:',ac)
		if iteration % 100 == 0:
			x_dt = mnist.test.images
			y_dt = mnist.test.labels
			acttl = 0
			for i in range(100):
				x_train = x_dt[i*100:i*100+100]
				y_train = y_dt[i*100:i*100+100]
				ac = sess.run(accuracy,feed_dict={inph:x_train, labh:y_train})
				acttl += ac
			print('Test acc:',acttl/100)
		if iteration % 400 == 0 and iteration>0:
			saver.save(sess,'./model/caps_model_iter%d.ckpt'%(iteration))


#variables
months = ["january", "february", "march", "april", "may", "june", "july", "august", "september", "october", "november", "december"]
february = "february"
months_30 = ["april", "june", "september", "november"]
months_31 = ["january", "march", "may", "july", "august", "october", "december"]

print ("Insert a name of month and I'll show you how many days there are in this:")
month = str(input())

#if the input is not a month's name there is an error
if month not in months:
    print ("Error! Insert a month's name:")
    month = str(input())

#calculating how many days there are in a month
if month == february:
    print ("The number of days in your month is 28/29 days, it depens if the year is leap or not.")

elif month in months_30:
    print ("The number of days in your month is 30 days!")

elif month in months_31:
    print ("The number of days in your month is 31 days!")


from urllib.request import urlretrieve
from os.path import isfile, isdir
from tqdm import tqdm
import tarfile

from tests import test_folder_path, test_cifar10_tar_data_exists

cifar10_dataset_folder_path = 'cifar-10-batches-py'
floyd_cifar10_location = '/input/cifar-10/python.tar.gz'
if isfile(floyd_cifar10_location):
    tar_gz_path = floyd_cifar10_location
else:
    tar_gz_path = 'cifar-10-python.tar.gz'

class DownloadProgress(tqdm):
    last_block = 0

    def hook(self, block_num=1, block_size=1, total_size=None):
        self.total = total_size
        self.update((block_num - self.last_block) * block_size)
        self.last_block = block_num

def download_cifar10_tar():
    if not isfile(tar_gz_path):
        with DownloadProgress(unit='B', unit_scale=True, miniters=1, desc='CIFAR-10 Dataset') as pbar:
            urlretrieve(
                'https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz',
                tar_gz_path,
                pbar.hook)

    test_cifar10_tar_data_exists(tar_gz_path)

def unzip():
    if not isdir(cifar10_dataset_folder_path):
        with tarfile.open(tar_gz_path) as tar:
            tar.extractall()
            tar.close()

def download_and_unzip():
    download_cifar10_tar()
    unzip()
    test_folder_path(cifar10_dataset_folder_path)


from django.http import JsonResponse
from django.shortcuts import render

from userman2.model import action

def displayActions(request):
    return render(request, "actions.html")

def getActions(request):
    dn = request.GET.get("dn", None)
    filter_data = False
    if dn is not None:
        filter_data = {"affectedDN": dn}
    actions = action.GetAllActions(filter_data)
    ret = []

    def transform_action(a):
        children = list(map(transform_action, filter(lambda x: x.parentDN == a.dn, actions)))
        return {
            "dn": a.dn,
            "affectedDN": a.affectedDN,
            "name": a.actionName,
            "description": a.description,
            "locked": a.locked,
            "host": a.host,
            "children": children,
        }

    for a in actions:
        if a.parentDN:
            continue
        ret.append(transform_action(a))
    return JsonResponse({"actions": ret})


from itertools import permutations

class Solution:
    def binaryGap(self, N):
        """
        :type N: int
        :rtype: int
        """
        maxGap = 0
        currGap = 0
        firstOne = False

        binString = "{0:b}".format(N)

        for bin in binString:
            if bin == "0" and firstOne:
                currGap += 1
            else:
                if not firstOne:
                    firstOne = True
                else:
                    maxGap = max(maxGap, currGap + 1)
                    currGap = 0

        return maxGap

    def reorderedPowerOf2(self, N):
        """
        :type N: int
        :rtype: bool
        """
        if N == 1:
            return True

        return self.permuteNum("", str(N))

    def permuteNum(self, prefix, str):
        n = len(str)
        if n == 0:
            return self.checkPowerOf(int(prefix))
        else:
            for i in range(n):
                self.permuteNum(prefix + str[i], str[0:i] + str[i+1:n])

    def checkPowerOf(self, a):
        return (a and (not(a & (a - 1))))


# import scipy.io.wavfile as wav
import speech_recognition as sr     # import the library

import uuid
import urllib
import os
from config  import configuration
import logging

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

class SpeechToText():

    def __init__(self, modelType="speechRecog"):

        self.staudio = configuration.staudio
        self.modelType = modelType
        if modelType =="deepspeech":
            # deepspeech model
            import deepspeech
            self.ds = self.load_deepspeech_model()
        elif modelType == "speechwavenet":
            # speech to text wavenet model
            from wavenet import recognize
        else:
            # speech recognition model
            self.speechRecog = sr.Recognizer()

    def load_deepspeech_model(self):
        N_FEATURES = 25
        N_CONTEXT = 9
        BEAM_WIDTH = 500
        LM_ALPHA = 0.75
        LM_BETA = 1.85
        ds = deepspeech.Model('deepspeech_model/deepspeech-0.7.3-models.pbmm')
        return ds

    def output_fileCreate(self, text):
        sid = str(uuid.uuid4())
        received_file = self.staudio + 'output_' + sid + '.wav'
        urllib.request.urlretrieve(text, received_file)
        # path = os.path.dirname(__file__)
        # return "audio/output_{0}.wav".format(sid)
        return received_file
        # return wav.read("audio/output_{0}.wav".format(sid))

    def predictText(self, text):
        audioFile = self.output_fileCreate(text)
        logger.debug("TS file name" + audioFile)
        message = ""
        if self.modelType =="deepspeech":
            audio, fs = wav.read(audioFile)
            message = self.ds.stt(audio, fs)
        elif self.modelType == "speechwavenet":
            message = recognize.predict_wavenet(audioFile)
        else:
            message = self.predict_recognition(audioFile)
        return {"text": message}

    # speerch recognition prediction
    def predict_recognition(self, filename):
        # recognize the chunk
        with sr.AudioFile(filename) as source:
            # remove this if it is not working
            # correctly.
            self.speechRecog.adjust_for_ambient_noise(source)
            audio_listened = self.speechRecog.listen(source)
        try:
            # try converting it to text
            # rec = r.recognize_sphinx(audio_listened)
            rec = self.speechRecog.recognize_google(audio_listened)
            # write the output to the file.
            return rec

            # catch any errors.
        except sr.UnknownValueError:
            msg = "Could not understand audio"
            return msg

        except sr.RequestError as e:
            msg = "Could not request results. check your internet connection"
            return msg


from setuptools import setup

with open("README.md", "r") as fh:
    long_description = fh.read()

setup(
    name="mreschke-serverbackups",
    version="1.0.3",
    description="Config driven server backups that feel like a simple backup.py script",
    long_description=long_description,
    long_description_content_type="text/markdown",

    author="Matthew Reschke",
    author_email="mail@mreschke.com",
    license="MIT",
    url="http://github.com/mreschke/server-backups",

    # I am using Native namespaces see https://packaging.python.org/guides/packaging-namespace-packages/#native-namespace-packages
    packages=["mreschke.serverbackups"],
    package_data={
        'mreschke.serverbackups': ['templates/*'],
    },
    zip_safe=False,
    python_requires=">=3.5",
    classifiers=[
        "Programming Language :: Python :: 3",
        "Programming Language :: Python :: 3.5",
        "Programming Language :: Python :: 3.6",
        "Programming Language :: Python :: 3.7",
        "Programming Language :: Python :: 3.8",
        "License :: OSI Approved :: MIT License",
        "Operating System :: Unix",
        "Environment :: Console",
    ],
    install_requires=[
        "click==7.1.*",
        "colored==1.4.*",
        "prettyprinter==0.18.*",
        "PyYAML==5.*",
    ],
    entry_points={
        "console_scripts": [
            "serverbackups=mreschke.serverbackups.init:cli"
        ],
    },
)


import os
import sys
import csv
import math
import subprocess

program_path = "bin/file_processor"
number_of_executions = 5

def main():

    file_sizes = [30, 1000, 10000, 100000, 1000000, 10000000, 100000000]

    min_times = {}
    max_times = {}

    for file_size in file_sizes:
        best_times = {}
        test_file_name = f"{file_size}.input"
        test_file_path = os.path.join("testes", f"{test_file_name}")
        verification_file_path = os.path.join("testes", f"{file_size}.txt")
        verification_data = get_verification_data(verification_file_path)

        buffer_size_start = 2 if file_size < 100000000 else 1024

        available_block_sizes = [i for i in powers_of_two_generator(file_size)]
        available_buffer_sizes = [i for i in powers_of_two_generator(file_size, buffer_size_start)]

        available_combinations = [[i, j] for i in available_block_sizes
                                         for j in available_buffer_sizes]

        for index, combination in enumerate(available_combinations):
            print(f"Running combination {index + 1} of {len(available_combinations)} for file: {test_file_name}...")

            block_size = combination[0]
            buffer_size = combination[1]

            best_time = get_best_execution_time_data(block_size, buffer_size, test_file_path, verification_data)

            best_times[(block_size, buffer_size)] = best_time

        min_times[file_size] = find_min_time(best_times)
        max_times[file_size] = find_max_time(best_times)

        output_best_times(best_times, file_size)

    output_summary(file_sizes, min_times, max_times)

def run_program_params(program_path, *params):
    return subprocess.run([program_path, *params], stdout=subprocess.PIPE).stdout.decode()

def confirm_response(program_response, verification_data):
    response_lines = program_response.splitlines()

    for index, line in enumerate(verification_data.splitlines()):
        if line != response_lines[index]:
            return False

    return True

def get_verification_data(verification_file_path):
    with open(verification_file_path, "r") as f:
            f.readline() # First line is discarded
            verification_data = f.read()

    return verification_data

def get_time(program_response):
    time_line = program_response.splitlines()[-1]
    time_str = time_line.replace("Time: ", "").replace("s", "")

    return float(time_str)

def get_best_execution_time_data(block_size, buffer_size, test_file_path, verification_data):
    times = []

    for _ in range(number_of_executions):
        program_response = run_program_params(program_path, f"{block_size}", f"{buffer_size}", test_file_path)

        is_correct = confirm_response(program_response, verification_data)

        if not is_correct:
            report_error_and_exit(program_response, verification_data, block_size, buffer_size)

        time = get_time(program_response)
        times.append(time)

    return min(times)

def find_min_time(best_times):
    current_min_key = None
    current_min_time = math.inf

    for key, time in best_times.items():
        if time < current_min_time:
            current_min_key = key
            current_min_time = time

    (block_size, buffer_size) = current_min_key
    return [block_size, buffer_size, current_min_time]

def find_max_time(best_times):
    current_max_key = None
    current_max_time = -math.inf

    for key, time in best_times.items():
        if time > current_max_time:
            current_max_key = key
            current_max_time = time

    (block_size, buffer_size) = current_max_key
    return [block_size, buffer_size, current_max_time]

def powers_of_two_generator(upper_bound, start=2):
    i = start
    while i <= upper_bound:
        yield i
        i *= 2

def report_error_and_exit(program_response, verification_data, block_size, buffer_size):
    print(f"Verification failed for block size: {block_size} and buffer size: {buffer_size}")
    print("\nExpected:")
    print(verification_data)
    print("\nGot:")
    print(program_response)

    sys.exit(1)

def output_best_times(best_times, file_size):
    output_directory = "results"
    output_file = os.path.join(output_directory, f"times_{file_size}.csv")

    if not os.path.exists(output_directory):
        os.makedirs(output_directory)

    best_times_csv = best_times_to_csv_format(best_times)

    with open(output_file, "w") as f:
        writer = csv.writer(f, dialect='excel')
        writer.writerows(best_times_csv)

def output_summary(file_sizes, min_times, max_times):
    output_directory = "results"
    output_file = os.path.join(output_directory, "summary.csv")

    if not os.path.exists(output_directory):
        os.makedirs(output_directory)

    summary_csv = summary_to_csv_format(file_sizes, min_times, max_times)

    with open(output_file, "w") as f:
        writer = csv.writer(f, dialect='excel')
        writer.writerows(summary_csv)

def best_times_to_csv_format(best_times):
    rows = [["block size", "buffer size", "time"]]

    for (block_size, buffer_size), time in best_times.items():
        rows.append([block_size, buffer_size, time])

    return rows

def summary_to_csv_format(file_sizes, min_times, max_times):
    rows = [["file size", "min block size", "min buffer size", "min time", "max block size", "max buffer size", "max time"]]

    for file_size in file_sizes:
        min_time = min_times[file_size]
        max_time = max_times[file_size]

        rows.append([file_size, *min_time, *max_time])

    return rows

if __name__ == "__main__":
    main()


import os, sys
import pandas as pd
import numpy as np
import json, copy
from .split_datasets import read_data,ALL_KEYS,ONE_KEYS,extend_multi_concepts,save_dcur
from .split_datasets import train_test_split,KFold_split,calStatistics,get_max_concepts,id_mapping,write_config

def generate_sequences(df, effective_keys, min_seq_len=3, maxlen = 200, pad_val = -1):
    save_keys = list(effective_keys) + ["selectmasks"]
    dres = {"selectmasks": []}
    dropnum = 0
    for i, row in df.iterrows():
        dcur = save_dcur(row, effective_keys)

        rest, lenrs = len(dcur["responses"]), len(dcur["responses"])
        j = 0
        while lenrs >= j + maxlen:
            rest = rest - (maxlen)
            for key in effective_keys:
                dres.setdefault(key, [])
                if key not in ONE_KEYS:
                    dres[key].append(",".join(dcur[key][j: j + maxlen]))#[str(k) for k in dcur[key][j: j + maxlen]]))
                else:
                    dres[key].append(dcur[key])
            dres["selectmasks"].append(",".join(["1"] * maxlen))

            j += maxlen
        if rest < min_seq_len:# delete sequence len less than min_seq_len
            dropnum += rest
            continue

        pad_dim = maxlen - rest
        for key in effective_keys:
            dres.setdefault(key, [])
            if key not in ONE_KEYS:
                paded_info = np.concatenate([dcur[key][j:], np.array([pad_val] * pad_dim)])
                dres[key].append(",".join([str(k) for k in paded_info]))
            else:
                dres[key].append(dcur[key])
        dres["selectmasks"].append(",".join(["1"] * rest + [str(pad_val)] * pad_dim))

    # after preprocess data, report
    dfinal = dict()
    for key in ALL_KEYS:
        if key in save_keys:
            dfinal[key] = dres[key]
    finaldf = pd.DataFrame(dfinal)
    print(f"dropnum: {dropnum}")
    return finaldf

def generate_window_sequences(df, effective_keys, maxlen=200, pad_val=-1):
    save_keys = list(effective_keys) + ["selectmasks"]
    dres = {"selectmasks": []}
    for i, row in df.iterrows():
        dcur = save_dcur(row, effective_keys)
        lenrs = len(dcur["responses"])
        if lenrs > maxlen:
            for key in effective_keys:
                dres.setdefault(key, [])
                if key not in ONE_KEYS:
                    dres[key].append(",".join(dcur[key][0: maxlen]))#[str(k) for k in dcur[key][0: maxlen]]))
                else:
                    dres[key].append(dcur[key])
            dres["selectmasks"].append(",".join(["1"] * maxlen))
            for j in range(maxlen+1, lenrs+1):
                for key in effective_keys:
                    dres.setdefault(key, [])
                    if key not in ONE_KEYS:
                        dres[key].append(",".join([str(k) for k in dcur[key][j-maxlen: j]]))
                    else:
                        dres[key].append(dcur[key])
                dres["selectmasks"].append(",".join([str(pad_val)] * (maxlen - 1) + ["1"]))
        else:
            for key in effective_keys:
                dres.setdefault(key, [])
                if key not in ONE_KEYS:
                    pad_dim = maxlen - lenrs
                    paded_info = np.concatenate([dcur[key][0:], np.array([pad_val] * pad_dim)])
                    dres[key].append(",".join([str(k) for k in paded_info]))
                else:
                    dres[key].append(dcur[key])
            dres["selectmasks"].append(",".join(["1"] * lenrs + [str(pad_val)] * pad_dim))

    dfinal = dict()
    for key in ALL_KEYS:
        if key in save_keys:
            # print(f"key: {key}, len: {len(dres[key])}")
            dfinal[key] = dres[key]
    finaldf = pd.DataFrame(dfinal)
    return finaldf

def save_id2idx(dkeyid2idx, save_path):
    with open(save_path, "w+") as fout:
        fout.write(json.dumps(dkeyid2idx))

def id_mapping_que(df):
    id_keys = ["questions", "concepts","uid"]
    dres = dict()
    dkeyid2idx = dict()
    print(f"df.columns: {df.columns}")
    for key in df.columns:
        if key not in id_keys:
            dres[key] = df[key]
    for i, row in df.iterrows():
        for key in id_keys:
            if key not in df.columns:
                continue
            dkeyid2idx.setdefault(key, dict())
            dres.setdefault(key, [])
            curids = []
            for id in row[key].split(","):
                sub_ids = id.split('_')
                sub_curids = []
                for sub_id in sub_ids:
                    if sub_id not in dkeyid2idx[key]:
                        dkeyid2idx[key][sub_id] = len(dkeyid2idx[key])
                    sub_curids.append(str(dkeyid2idx[key][sub_id]))
                curids.append("_".join(sub_curids))
            dres[key].append(",".join(curids))
    finaldf = pd.DataFrame(dres)
    return finaldf, dkeyid2idx

def main(dname, fname, dataset_name, configf, min_seq_len = 3, maxlen = 200, kfold = 5):
    """split main function

    Args:
        dname (str): data folder path
        fname (str): the data file used to split, needs 6 columns, format is: (NA indicates the dataset has no corresponding info)
            uid,seqlen: 50121,4
            quetion ids: NA
            concept ids: 7014,7014,7014,7014
            responses: 0,1,1,1
            timestamps: NA
            cost times: NA
        dataset_name (str): dataset name
        configf (str): the dataconfig file path
        min_seq_len (int, optional): the min seqlen, sequences less than this value will be filtered out. Defaults to 3.
        maxlen (int, optional): the max seqlen. Defaults to 200.
        kfold (int, optional): the folds num needs to split. Defaults to 5.

    """
    stares = []

    total_df, effective_keys = read_data(fname)
    #cal max_concepts
    if 'concepts' in effective_keys:
        max_concepts = get_max_concepts(total_df)
    else:
        max_concepts = -1

    oris, _, qs, cs, seqnum = calStatistics(total_df, stares, "original")
    print("="*20)
    print(f"original total interactions: {oris}, qs: {qs}, cs: {cs}, seqnum: {seqnum}")

     # just for id map
    total_df, dkeyid2idx = id_mapping_que(total_df)
    dkeyid2idx["max_concepts"] = max_concepts

    save_id2idx(dkeyid2idx, os.path.join(dname, "keyid2idx.json"))
    effective_keys.add("fold")

    df_save_keys = []
    for key in ALL_KEYS:
        if key in effective_keys:
            df_save_keys.append(key)

    # train test split
    train_df, test_df = train_test_split(total_df, 0.2)
    splitdf = KFold_split(train_df, kfold)
    splitdf[df_save_keys].to_csv(os.path.join(dname, "train_valid_quelevel.csv"), index=None)
    ins, ss, qs, cs, seqnum = calStatistics(splitdf, stares, "original train+valid question level")
    print(f"train+valid original interactions num: {ins}, select num: {ss}, qs: {qs}, cs: {cs}, seqnum: {seqnum}")

    # generate sequences
    split_seqs = generate_sequences(splitdf, effective_keys, min_seq_len, maxlen)
    ins, ss, qs, cs, seqnum = calStatistics(split_seqs, stares, "train+valid sequences question level")
    print(f"train+valid sequences interactions num: {ins}, select num: {ss}, qs: {qs}, cs: {cs}, seqnum: {seqnum}")
    split_seqs.to_csv(os.path.join(dname, "train_valid_sequences_quelevel.csv"), index=None)
    # print(f"split seqs dtypes: {split_seqs.dtypes}")

    # for test dataset
    # add default fold -1 to test!
    test_df["fold"] = [-1] * test_df.shape[0]
    test_seqs = generate_sequences(test_df, list(effective_keys), min_seq_len, maxlen)
    ins, ss, qs, cs, seqnum = calStatistics(test_df, stares, "test original question level")
    print(f"original test interactions num: {ins}, select num: {ss}, qs: {qs}, cs: {cs}, seqnum: {seqnum}")
    ins, ss, qs, cs, seqnum = calStatistics(test_seqs, stares, "test sequences question level")
    print(f"test sequences interactions num: {ins}, select num: {ss}, qs: {qs}, cs: {cs}, seqnum: {seqnum}")
    print("="*20)

    test_window_seqs = generate_window_sequences(test_df, list(effective_keys), maxlen)

    test_df = test_df[df_save_keys]
    test_df.to_csv(os.path.join(dname, "test_quelevel.csv"), index=None)
    test_seqs.to_csv(os.path.join(dname, "test_sequences_quelevel.csv"), index=None)
    test_window_seqs.to_csv(os.path.join(dname, "test_window_sequences_quelevel.csv"), index=None)

    ins, ss, qs, cs, seqnum = calStatistics(test_window_seqs, stares, "test window question level")
    print(f"test window interactions num: {ins}, select num: {ss}, qs: {qs}, cs: {cs}, seqnum: {seqnum}")

    other_config = {
        "train_valid_original_file_quelevel": "train_valid_quelevel.csv",
        "train_valid_file_quelevel": "train_valid_sequences_quelevel.csv",
        "test_file_quelevel": "test_sequences_quelevel.csv",
        "test_window_file_quelevel": "test_window_sequences_quelevel.csv",
        "test_original_file_quelevel": "test_quelevel.csv"
    }

    write_config(dataset_name=dataset_name, dkeyid2idx=dkeyid2idx, effective_keys=effective_keys,
                configf=configf, dpath = dname, k=kfold,min_seq_len = min_seq_len, maxlen=maxlen,other_config=other_config)

    print("="*20)
    print("\n".join(stares))


from swagger_tester import swagger_test

# Dict containing the error you don't want to raise.
# By default, every status_code over other than 1xx, 2xx or 3xx
# will be considered as an error.
authorize_error = {
    'get': {
        '/artifact/{storage_backend}/{bucket_name}/{artifact_name}': [200]
    }
}

# Run the test with connexion
# An AssertionError will be raise in case of error.
swagger_test('/app/storage_proxy/swagger.yaml', authorize_error=authorize_error)


import torch
from torch import nn
import torch.nn.functional as F
import math
import copy
from util.box_ops import box_cxcywh_to_xyxy, generalized_box_iou
from util.misc import (NestedTensor, nested_tensor_from_tensor_list,
                       accuracy, get_world_size, interpolate,
                       is_dist_avail_and_initialized, inverse_sigmoid)
import numpy as np
import clip
from datasets.hico_text_label import hico_text_label, hico_obj_text_label, hico_unseen_index
from datasets.vcoco_text_label import vcoco_hoi_text_label, vcoco_obj_text_label

from .backbone import build_backbone
from .matcher import build_matcher
from .deformable_transformer import build_deformable_transformer
from .utils import sigmoid_focal_loss, MLP

import sys
from omegaconf import OmegaConf
from einops import rearrange, repeat

from ldm.util import instantiate_from_config
from SD_Extractor import UNetWrapper, TextAdapter

def _sigmoid(x):
    y = torch.clamp(x.sigmoid(), min=1e-4, max=1 - 1e-4)
    return y

class DiffHOI_L(nn.Module):
    """ This is the Cross-Attention Detector module that performs object detection """

    def __init__(self, backbone, transformer, num_classes, num_queries,
                 aux_loss=False, iter_update=False,
                 query_dim=2,
                 random_refpoints_xy=False,
                 fix_refpoints_hw=-1,
                 num_feature_levels=1,
                 nheads=8,
                 # two stage
                 two_stage_type='no',  # ['no', 'standard']
                 two_stage_add_query_num=0,
                 dec_pred_class_embed_share=True,
                 dec_pred_bbox_embed_share=True,
                 two_stage_class_embed_share=True,
                 two_stage_bbox_embed_share=True,
                 decoder_sa_type='sa',
                 num_patterns=0,
                 dn_number=100,
                 dn_box_noise_scale=0.4,
                 dn_label_noise_ratio=0.5,
                 dn_labelbook_size=100,
                 args=None,
                 unet_config = dict()
                 ):
        """ Initializes the model.
        Parameters:
            backbone: torch module of the backbone to be used. See backbone.py
            transformer: torch module of the transformer architecture. See transformer.py
            num_classes: number of object classes
            num_queries: number of object queries, ie detection slot. This is the maximal number of objects
                         Conditional DETR can detect in a single image. For COCO, we recommend 100 queries.
            aux_loss: True if auxiliary decoding losses (loss at each decoder layer) are to be used.

            fix_refpoints_hw: -1(default): learn w and h for each box seperately
                                >0 : given fixed number
                                -2 : learn a shared w and h
        """
        super().__init__()

        config_path = os.environ.get("SD_Config")
        ckpt_path = os.environ.get("SD_ckpt")
        config = OmegaConf.load(config_path)
        config.model.params.ckpt_path = ckpt_path
        config.model.params.cond_stage_config.target = 'ldm.modules.encoders.modules.AbstractEncoder'
        sd_model = instantiate_from_config(config.model)
        self.encoder_vq = sd_model.first_stage_model
        self.unet = UNetWrapper(sd_model.model, **unet_config)
        device = "cuda" if torch.cuda.is_available() else "cpu"
        self.unet.to(device)
        # self.unet.eval()
        sd_model.model = None
        sd_model.first_stage_model = None
        del sd_model.cond_stage_model
        del self.encoder_vq.decoder
        self.sd_model = sd_model
        self.text_adapter = nn.Linear(args.clip_embed_dim, 768)
        self.clip_adapter = nn.Linear(args.clip_embed_dim,args.clip_embed_dim)

        self.num_queries = num_queries
        self.transformer = transformer
        self.num_classes = num_classes
        self.hidden_dim = hidden_dim = transformer.d_model
        self.num_feature_levels = num_feature_levels
        self.nheads = nheads
        self.label_enc = nn.Embedding(dn_labelbook_size + 1, hidden_dim)
        self.args=args
        # setting query dim
        self.query_dim = query_dim
        assert query_dim == 4
        self.random_refpoints_xy = random_refpoints_xy
        self.fix_refpoints_hw = fix_refpoints_hw

        # for dn training
        self.num_patterns = num_patterns
        self.dn_number = dn_number
        self.dn_box_noise_scale = dn_box_noise_scale
        self.dn_label_noise_ratio = dn_label_noise_ratio
        self.dn_labelbook_size = dn_labelbook_size

        # prepare input projection layers
        if num_feature_levels > 1:
            num_backbone_outs = len(backbone.num_channels)
            input_proj_list = []
            for _ in range(num_backbone_outs):
                in_channels = backbone.num_channels[_]
                input_proj_list.append(nn.Sequential(
                    nn.Conv2d(in_channels, hidden_dim, kernel_size=1),
                    nn.GroupNorm(32, hidden_dim),
                ))
            for _ in range(num_feature_levels - num_backbone_outs):
                input_proj_list.append(nn.Sequential(
                    nn.Conv2d(in_channels, hidden_dim, kernel_size=3, stride=2, padding=1),
                    nn.GroupNorm(32, hidden_dim),
                ))
                in_channels = hidden_dim
            self.input_proj = nn.ModuleList(input_proj_list)
        else:
            assert two_stage_type == 'no', "two_stage_type should be no if num_feature_levels=1 !!!"
            self.input_proj = nn.ModuleList([
                nn.Sequential(
                    nn.Conv2d(backbone.num_channels[-1], hidden_dim, kernel_size=1),
                    nn.GroupNorm(32, hidden_dim),
                )])

        self.backbone = backbone
        self.aux_loss = aux_loss
        self.box_pred_damping = box_pred_damping = None

        self.iter_update = iter_update
        assert iter_update, "Why not iter_update?"

        # prepare pred layers
        self.dec_pred_class_embed_share = dec_pred_class_embed_share
        self.dec_pred_bbox_embed_share = dec_pred_bbox_embed_share
        # prepare class & box embed
        _class_embed = nn.Linear(hidden_dim, num_classes)
        _bbox_embed = MLP(hidden_dim, hidden_dim, 4, 3)

        # init the two embed layers
        prior_prob = 0.01
        bias_value = -math.log((1 - prior_prob) / prior_prob)
        _class_embed.bias.data = torch.ones(self.num_classes) * bias_value
        nn.init.constant_(_bbox_embed.layers[-1].weight.data, 0)
        nn.init.constant_(_bbox_embed.layers[-1].bias.data, 0)

        if dec_pred_bbox_embed_share:
            box_embed_layerlist = [_bbox_embed for i in range(transformer.num_decoder_layers -3 )]
        else:
            box_embed_layerlist = [copy.deepcopy(_bbox_embed) for i in range(transformer.num_decoder_layers -3  )]
        if dec_pred_class_embed_share:
            class_embed_layerlist = [_class_embed for i in range(transformer.num_decoder_layers -3  )]
        else:
            class_embed_layerlist = [copy.deepcopy(_class_embed) for i in range(transformer.num_decoder_layers -3 )]

        self.bbox_embed = nn.ModuleList(box_embed_layerlist)
        self.class_embed = nn.ModuleList(class_embed_layerlist)
        self.transformer.decoder.bbox_embed = self.bbox_embed
        self.transformer.decoder.class_embed = self.class_embed

        # two stage
        self.two_stage_type = two_stage_type
        self.two_stage_add_query_num = two_stage_add_query_num
        assert two_stage_type in ['no', 'standard'], "unknown param {} of two_stage_type".format(two_stage_type)
        if two_stage_type != 'no':
            if two_stage_bbox_embed_share:
                assert dec_pred_class_embed_share and dec_pred_bbox_embed_share
                self.transformer.enc_out_bbox_embed = _bbox_embed
            else:
                self.transformer.enc_out_bbox_embed = copy.deepcopy(_bbox_embed)

            if two_stage_class_embed_share:
                assert dec_pred_class_embed_share and dec_pred_bbox_embed_share
                self.transformer.enc_out_class_embed = _class_embed
            else:
                self.transformer.enc_out_class_embed = copy.deepcopy(_class_embed)

            self.refpoint_embed = None
            if self.two_stage_add_query_num > 0:
                self.init_ref_points(two_stage_add_query_num)

        self.decoder_sa_type = decoder_sa_type
        assert decoder_sa_type in ['sa', 'ca_label', 'ca_content']
        if decoder_sa_type == 'ca_label':
            self.label_embedding = nn.Embedding(num_classes, hidden_dim)
            for layer in self.transformer.decoder.layers:
                layer.label_embedding = self.label_embedding
        else:
            for layer in self.transformer.decoder.layers:
                layer.label_embedding = None
            self.label_embedding = None

        self.dec_layers = 3

        self._reset_parameters()

        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
        self.obj_logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))

        if self.args.dataset_file == 'hico':
            hoi_text_label = hico_text_label
            obj_text_label = hico_obj_text_label
            unseen_index = hico_unseen_index
        elif self.args.dataset_file == 'vcoco':
            hoi_text_label = vcoco_hoi_text_label
            obj_text_label = vcoco_obj_text_label
            unseen_index = None

        clip_label, obj_clip_label, v_linear_proj_weight, hoi_text, obj_text, train_clip_label = \
            self.init_classifier_with_CLIP(hoi_text_label, obj_text_label, unseen_index)
        num_obj_classes = len(obj_text) - 1  # del nothing

        self.hoi_class_fc = nn.Sequential(
            nn.Linear(hidden_dim, args.clip_embed_dim),
            nn.LayerNorm(args.clip_embed_dim),
        )

        if args.with_clip_label:
            self.visual_projection = nn.Linear(args.clip_embed_dim, len(hoi_text))
            self.visual_projection.weight.data = train_clip_label / train_clip_label.norm(dim=-1, keepdim=True)
            if self.args.dataset_file == 'hico' and self.args.zero_shot_type != 'default':
                self.eval_visual_projection = nn.Linear(args.clip_embed_dim, 600)
                self.eval_visual_projection.weight.data = clip_label / clip_label.norm(dim=-1, keepdim=True)

        if args.with_obj_clip_label:
            self.obj_class_fc = nn.Sequential(
                nn.Linear(hidden_dim, args.clip_embed_dim),
                nn.LayerNorm(args.clip_embed_dim),
            )
            self.obj_visual_projection = nn.Linear(args.clip_embed_dim, num_obj_classes + 1)
            self.obj_visual_projection.weight.data = obj_clip_label / obj_clip_label.norm(dim=-1, keepdim=True)
        device = "cuda" if torch.cuda.is_available() else "cpu"
        self.clip_model, _ = clip.load(args.clip_model, device=device)

        self.image_adapter_4 = nn.Conv2d(1280, hidden_dim, kernel_size=1)

        self.image_adapter_3 = nn.Conv2d(1280, hidden_dim, kernel_size=1)

        self.image_adapter_2 = nn.Conv2d(640, hidden_dim, kernel_size=1)

        self.image_adapter_1 = nn.Conv2d(320, hidden_dim, kernel_size=1)

        self.clip_adapter = nn.Linear(args.clip_embed_dim,args.clip_embed_dim)
        # import pdb
        # pdb.set_trace()

        self.freeze()

    def freeze(self):
        for param in self.class_embed.parameters():
            param.requires_grad = False

        for param in self.transformer.enc_out_bbox_embed.parameters():
            param.requires_grad = False

    def _reset_parameters(self):
        # init input_proj
        for proj in self.input_proj:
            nn.init.xavier_uniform_(proj[0].weight, gain=1)
            nn.init.constant_(proj[0].bias, 0)

    def init_classifier_with_CLIP(self, hoi_text_label, obj_text_label, unseen_index):
        device = "cuda" if torch.cuda.is_available() else "cpu"
        text_inputs = torch.cat([clip.tokenize(hoi_text_label[id]) for id in hoi_text_label.keys()])
        if self.args.del_unseen and unseen_index is not None:
            hoi_text_label_del = {}
            unseen_index_list = unseen_index.get(self.args.zero_shot_type, [])
            for idx, k in enumerate(hoi_text_label.keys()):
                if idx in unseen_index_list:
                    continue
                else:
                    hoi_text_label_del[k] = hoi_text_label[k]
        else:
            hoi_text_label_del = hoi_text_label.copy()
        text_inputs_del = torch.cat(
            [clip.tokenize(hoi_text_label[id]) for id in hoi_text_label_del.keys()])

        obj_text_inputs = torch.cat([clip.tokenize(obj_text[1]) for obj_text in obj_text_label])
        clip_model, preprocess = clip.load(self.args.clip_model, device=device)
        with torch.no_grad():
            text_embedding = clip_model.encode_text(text_inputs.to(device))
            text_embedding_del = clip_model.encode_text(text_inputs_del.to(device))
            obj_text_embedding = clip_model.encode_text(obj_text_inputs.to(device))
            v_linear_proj_weight = clip_model.visual.proj.detach()

        del clip_model

        return text_embedding.float(), obj_text_embedding.float(), v_linear_proj_weight.float(), \
               hoi_text_label_del, obj_text_inputs, text_embedding_del.float()

    def init_ref_points(self, use_num_queries):
        self.refpoint_embed = nn.Embedding(use_num_queries, self.query_dim)
        if self.random_refpoints_xy:
            self.refpoint_embed.weight.data[:, :2].uniform_(0, 1)
            self.refpoint_embed.weight.data[:, :2] = inverse_sigmoid(self.refpoint_embed.weight.data[:, :2])
            self.refpoint_embed.weight.data[:, :2].requires_grad = False

        if self.fix_refpoints_hw > 0:
            print("fix_refpoints_hw: {}".format(self.fix_refpoints_hw))
            assert self.random_refpoints_xy
            self.refpoint_embed.weight.data[:, 2:] = self.fix_refpoints_hw
            self.refpoint_embed.weight.data[:, 2:] = inverse_sigmoid(self.refpoint_embed.weight.data[:, 2:])
            self.refpoint_embed.weight.data[:, 2:].requires_grad = False
        elif int(self.fix_refpoints_hw) == -1:
            pass
        elif int(self.fix_refpoints_hw) == -2:
            print('learn a shared h and w')
            assert self.random_refpoints_xy
            self.refpoint_embed = nn.Embedding(use_num_queries, 2)
            self.refpoint_embed.weight.data[:, :2].uniform_(0, 1)
            self.refpoint_embed.weight.data[:, :2] = inverse_sigmoid(self.refpoint_embed.weight.data[:, :2])
            self.refpoint_embed.weight.data[:, :2].requires_grad = False
            self.hw_embed = nn.Embedding(1, 1)
        else:
            raise NotImplementedError('Unknown fix_refpoints_hw {}'.format(self.fix_refpoints_hw))

    def extract_feat(self, img,targets):
        """Extract features from images."""
        target_clip_inputs = torch.cat([t['clip_inputs'].unsqueeze(0) for t in targets])
        target_sd_inputs = torch.cat([t['sd_inputs'].unsqueeze(0) for t in targets])
        with torch.no_grad():
            clip_feats = self.clip_model.encode_image(target_clip_inputs)
        clip_feats=self.text_adapter(clip_feats.unsqueeze(1).float())
        t = torch.zeros((img.shape[0],), device=img.device).long()
        with torch.no_grad():
            latents = self.encoder_vq.encode(target_sd_inputs)
            latents = latents.mode().detach()
            outs = self.unet(latents, t, c_crossattn=[clip_feats])
        return outs

    def forward(self, samples: NestedTensor, targets , is_training=True):
        """ The forward expects a NestedTensor, which consists of:
               - samples.tensor: batched images, of shape [batch_size x 3 x H x W]
               - samples.mask: a binary mask of shape [batch_size x H x W], containing 1 on padded pixels

            It returns a dict with the following elements:
               - "pred_logits": the classification logits (including no-object) for all queries.
                                Shape= [batch_size x num_queries x num_classes]
               - "pred_boxes": The normalized boxes coordinates for all queries, represented as
                               (center_x, center_y, width, height). These values are normalized in [0, 1],
                               relative to the size of each individual image (disregarding possible padding).
                               See PostProcess for information on how to retrieve the unnormalized bounding box.
               - "aux_outputs": Optional, only returned when auxilary losses are activated. It is a list of
                                dictionnaries containing the two above keys for each decoder layer.
        """
        if isinstance(samples, (list, torch.Tensor)):
            samples = nested_tensor_from_tensor_list(samples)
        features, poss = self.backbone(samples)

        srcs = []
        masks = []
        for l, feat in enumerate(features):
            src, mask = feat.decompose()
            srcs.append(self.input_proj[l](src))
            masks.append(mask)
            assert mask is not None
        if self.num_feature_levels > len(srcs):
            _len_srcs = len(srcs)
            for l in range(_len_srcs, self.num_feature_levels):
                if l == _len_srcs:
                    src = self.input_proj[l](features[-1].tensors)
                else:
                    src = self.input_proj[l](srcs[-1])
                m = samples.mask
                mask = F.interpolate(m[None].float(), size=src.shape[-2:]).to(torch.bool)[0]
                pos_l = self.backbone[1](NestedTensor(src, mask)).to(src.dtype)
                srcs.append(src)
                masks.append(mask)
                poss.append(pos_l)

        input_query_bbox = input_query_label = attn_mask = dn_meta = None

        sd_feat = self.extract_feat(samples.decompose()[0],targets)
        sd_decoder_4 = self.image_adapter_4(sd_feat[-1])
        sd_decoder_3 = self.image_adapter_3(sd_feat[-2])
        sd_decoder_2 = self.image_adapter_2(sd_feat[-3])
        sd_decoder_1 = self.image_adapter_1(sd_feat[-4])
        sd_decoder_4 = torch.nn.functional.interpolate(sd_decoder_4,size=[srcs[-1].shape[-2],srcs[-1].shape[-1]])
        sd_decoder_3 = torch.nn.functional.interpolate(sd_decoder_3,size=[srcs[-2].shape[-2],srcs[-2].shape[-1]])
        sd_decoder_2 = torch.nn.functional.interpolate(sd_decoder_2,size=[srcs[-3].shape[-2],srcs[-3].shape[-1]])
        sd_decoder_1 = torch.nn.functional.interpolate(sd_decoder_1,size=[srcs[-4].shape[-2],srcs[-4].shape[-1]])
        sd_decoder = [sd_decoder_1,sd_decoder_2,sd_decoder_3,sd_decoder_4]

        hs, hs_interaction , reference, hs_enc, ref_enc, init_box_proposal = self.transformer(srcs, masks, input_query_bbox, poss,
                                                                             input_query_label,sd_decoder, attn_mask)
        # In case num object=0
        hs[0] += self.label_enc.weight[0, 0] * 0.0

        hs_interaction = torch.stack(hs_interaction)
        # deformable-detr-like anchor update
        # reference_before_sigmoid = inverse_sigmoid(reference[:-1]) # n_dec, bs, nq, 4
        outputs_coord_list = []
        for dec_lid, (layer_ref_sig, layer_bbox_embed, layer_hs) in enumerate(zip(reference[:-1], self.bbox_embed, hs)):
            layer_delta_unsig = layer_bbox_embed(layer_hs)
            layer_outputs_unsig = layer_delta_unsig + inverse_sigmoid(layer_ref_sig)
            layer_outputs_unsig = layer_outputs_unsig.sigmoid()
            outputs_coord_list.append(layer_outputs_unsig)
        outputs_coord_list = torch.stack(outputs_coord_list)

        # clip_feat
        target_clip_inputs = torch.cat([t['clip_inputs'].unsqueeze(0) for t in targets])
        with torch.no_grad():
            target_clip_feats = self.clip_model.encode_image(target_clip_inputs)
        target_clip_feats = self.clip_adapter(target_clip_feats.float())

        if self.args.with_obj_clip_label:
            obj_logit_scale = self.obj_logit_scale.exp()
            o_hs = self.obj_class_fc(torch.stack(hs))
            o_hs = o_hs / o_hs.norm(dim=-1, keepdim=True)
            outputs_obj_class = obj_logit_scale * self.obj_visual_projection(o_hs)

        if self.args.with_clip_label:
            logit_scale = self.logit_scale.exp()
            inter_hs = self.hoi_class_fc(hs_interaction)
            inter_hs = inter_hs + target_clip_feats[None,:,None,:].repeat(inter_hs.shape[0],1,inter_hs.shape[2],1)
            inter_hs = inter_hs / inter_hs.norm(dim=-1, keepdim=True)
            if self.args.dataset_file == 'hico' and self.args.zero_shot_type != 'default' \
                    and (self.args.eval or not is_training):
                outputs_hoi_class = logit_scale * self.eval_visual_projection(inter_hs)
            else:
                outputs_hoi_class = logit_scale * self.visual_projection(inter_hs)

        outputs_coord_list = outputs_coord_list.view(outputs_coord_list.shape[0],outputs_coord_list.shape[1],int(self.num_queries/2),2,4)

        outputs_class = outputs_obj_class.view(outputs_obj_class.shape[0],outputs_obj_class.shape[1],int(self.num_queries/2),2,outputs_obj_class.shape[-1])

        out = {'pred_hoi_logits': outputs_hoi_class[-1], 'pred_obj_logits': outputs_class[:,:,:,1,:][-1],
               'pred_sub_boxes': outputs_coord_list[:,:,:,0,:][-1], 'pred_obj_boxes': outputs_coord_list[:,:,:,1,:][-1],}

        if self.aux_loss:
            out['aux_outputs'] = self._set_aux_loss_triplet(outputs_hoi_class, outputs_class[:,:,:,1,:],
                                                            outputs_coord_list[:,:,:,0,:], outputs_coord_list[:,:,:,1,:])

        out['dn_meta'] = dn_meta

        return out

    @torch.jit.unused
    def _set_aux_loss(self, outputs_class, outputs_coord):
        # this is a workaround to make torchscript happy, as torchscript
        # doesn't support dictionary with non-homogeneous values, such
        # as a dict having both a Tensor and a list.
        return [{'pred_logits': a, 'pred_boxes': b}
                for a, b in zip(outputs_class[:-1], outputs_coord[:-1])]

    @torch.jit.unused
    def _set_aux_loss_triplet(self, outputs_hoi_class, outputs_obj_class,
                              outputs_sub_coord, outputs_obj_coord, outputs_inter_hs=None):

        aux_outputs = {'pred_hoi_logits': outputs_hoi_class[-self.dec_layers: -1],
                       'pred_obj_logits': outputs_obj_class[-self.dec_layers: -1],
                       'pred_sub_boxes': outputs_sub_coord[-self.dec_layers: -1],
                       'pred_obj_boxes': outputs_obj_coord[-self.dec_layers: -1]}
        if outputs_inter_hs is not None:
            aux_outputs['inter_memory'] = outputs_inter_hs[-self.dec_layers: -1]
        outputs_auxes = []
        for i in range(self.dec_layers - 1):
            output_aux = {}
            for aux_key in aux_outputs.keys():
                output_aux[aux_key] = aux_outputs[aux_key][i]
            outputs_auxes.append(output_aux)
        return outputs_auxes

class MLP(nn.Module):
    """ Very simple multi-layer perceptron (also called FFN)"""

    def __init__(self, input_dim, hidden_dim, output_dim, num_layers):
        super().__init__()
        self.num_layers = num_layers
        h = [hidden_dim] * (num_layers - 1)
        self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]))

    def forward(self, x):
        for i, layer in enumerate(self.layers):
            x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
        return x

class SetCriterionHOI(nn.Module):

    def __init__(self, num_obj_classes, num_queries, num_verb_classes, matcher, weight_dict, eos_coef, losses, args):
        super().__init__()

        self.num_obj_classes = num_obj_classes
        self.num_queries = num_queries
        self.num_verb_classes = num_verb_classes
        self.matcher = matcher
        self.weight_dict = weight_dict
        self.eos_coef = eos_coef
        self.losses = losses
        empty_weight = torch.ones(self.num_obj_classes + 1)
        empty_weight[-1] = self.eos_coef
        self.register_buffer('empty_weight', empty_weight)
        device = "cuda" if torch.cuda.is_available() else "cpu"
        self.clip_model, _ = clip.load(args.clip_model, device=device)
        self.alpha = args.alpha

    def loss_obj_labels(self, outputs, targets, indices, num_interactions, log=True):
        assert 'pred_obj_logits' in outputs
        src_logits = outputs['pred_obj_logits']

        idx = self._get_src_permutation_idx(indices)
        target_classes_o = torch.cat([t['obj_labels'][J] for t, (_, J) in zip(targets, indices)])
        target_classes = torch.full(src_logits.shape[:2], self.num_obj_classes,
                                    dtype=torch.int64, device=src_logits.device)
        target_classes[idx] = target_classes_o

        loss_obj_ce = F.cross_entropy(src_logits.transpose(1, 2), target_classes, self.empty_weight)
        losses = {'loss_obj_ce': loss_obj_ce}

        if log:
            losses['obj_class_error'] = 100 - accuracy(src_logits[idx], target_classes_o)[0]
        return losses

    @torch.no_grad()
    def loss_obj_cardinality(self, outputs, targets, indices, num_interactions):
        pred_logits = outputs['pred_obj_logits']
        device = pred_logits.device
        tgt_lengths = torch.as_tensor([len(v['obj_labels']) for v in targets], device=device)
        card_pred = (pred_logits.argmax(-1) != pred_logits.shape[-1] - 1).sum(1)
        card_err = F.l1_loss(card_pred.float(), tgt_lengths.float())
        losses = {'obj_cardinality_error': card_err}
        return losses

    def loss_verb_labels(self, outputs, targets, indices, num_interactions):
        assert 'pred_verb_logits' in outputs
        src_logits = outputs['pred_verb_logits']

        idx = self._get_src_permutation_idx(indices)
        target_classes_o = torch.cat([t['verb_labels'][J] for t, (_, J) in zip(targets, indices)])
        target_classes = torch.zeros_like(src_logits)
        target_classes[idx] = target_classes_o

        src_logits = src_logits.sigmoid()
        loss_verb_ce = self._neg_loss(src_logits, target_classes, weights=None, alpha=self.alpha)
        losses = {'loss_verb_ce': loss_verb_ce}
        return losses

    def loss_hoi_labels(self, outputs, targets, indices, num_interactions, topk=5):
        assert 'pred_hoi_logits' in outputs
        src_logits = outputs['pred_hoi_logits']

        idx = self._get_src_permutation_idx(indices)
        target_classes_o = torch.cat([t['hoi_labels'][J] for t, (_, J) in zip(targets, indices)])
        target_classes = torch.zeros_like(src_logits)
        target_classes[idx] = target_classes_o
        src_logits = _sigmoid(src_logits)
        loss_hoi_ce = self._neg_loss(src_logits, target_classes, weights=None, alpha=self.alpha)
        losses = {'loss_hoi_labels': loss_hoi_ce}

        _, pred = src_logits[idx].topk(topk, 1, True, True)
        acc = 0.0
        for tid, target in enumerate(target_classes_o):
            tgt_idx = torch.where(target == 1)[0]
            if len(tgt_idx) == 0:
                continue
            acc_pred = 0.0
            for tgt_rel in tgt_idx:
                acc_pred += (tgt_rel in pred[tid])
            acc += acc_pred / len(tgt_idx)
        rel_labels_error = 100 - 100 * acc / max(len(target_classes_o), 1)
        losses['hoi_class_error'] = torch.from_numpy(np.array(
            rel_labels_error)).to(src_logits.device).float()
        return losses

    def loss_sub_obj_boxes(self, outputs, targets, indices, num_interactions):
        assert 'pred_sub_boxes' in outputs and 'pred_obj_boxes' in outputs
        idx = self._get_src_permutation_idx(indices)
        src_sub_boxes = outputs['pred_sub_boxes'][idx]
        src_obj_boxes = outputs['pred_obj_boxes'][idx]
        target_sub_boxes = torch.cat([t['sub_boxes'][i] for t, (_, i) in zip(targets, indices)], dim=0)
        target_obj_boxes = torch.cat([t['obj_boxes'][i] for t, (_, i) in zip(targets, indices)], dim=0)

        exist_obj_boxes = (target_obj_boxes != 0).any(dim=1)

        losses = {}
        if src_sub_boxes.shape[0] == 0:
            losses['loss_sub_bbox'] = src_sub_boxes.sum()
            losses['loss_obj_bbox'] = src_obj_boxes.sum()
            losses['loss_sub_giou'] = src_sub_boxes.sum()
            losses['loss_obj_giou'] = src_obj_boxes.sum()
        else:
            loss_sub_bbox = F.l1_loss(src_sub_boxes, target_sub_boxes, reduction='none')
            loss_obj_bbox = F.l1_loss(src_obj_boxes, target_obj_boxes, reduction='none')
            losses['loss_sub_bbox'] = loss_sub_bbox.sum() / num_interactions
            losses['loss_obj_bbox'] = (loss_obj_bbox * exist_obj_boxes.unsqueeze(1)).sum() / (
                    exist_obj_boxes.sum() + 1e-4)
            loss_sub_giou = 1 - torch.diag(generalized_box_iou(box_cxcywh_to_xyxy(src_sub_boxes),
                                                               box_cxcywh_to_xyxy(target_sub_boxes)))
            loss_obj_giou = 1 - torch.diag(generalized_box_iou(box_cxcywh_to_xyxy(src_obj_boxes),
                                                               box_cxcywh_to_xyxy(target_obj_boxes)))
            losses['loss_sub_giou'] = loss_sub_giou.sum() / num_interactions
            losses['loss_obj_giou'] = (loss_obj_giou * exist_obj_boxes).sum() / (exist_obj_boxes.sum() + 1e-4)
        return losses

    def mimic_loss(self, outputs, targets, indices, num_interactions):
        src_feats = outputs['inter_memory']
        src_feats = torch.mean(src_feats, dim=1)

        target_clip_inputs = torch.cat([t['clip_inputs'].unsqueeze(0) for t in targets])
        with torch.no_grad():
            target_clip_feats = self.clip_model.encode_image(target_clip_inputs)
        loss_feat_mimic = F.l1_loss(src_feats, target_clip_feats)
        losses = {'loss_feat_mimic': loss_feat_mimic}
        return losses

    def _neg_loss(self, pred, gt, weights=None, alpha=0.25):
        ''' Modified focal loss. Exactly the same as CornerNet.
          Runs faster and costs a little bit more memory
        '''
        pos_inds = gt.eq(1).float()
        neg_inds = gt.lt(1).float()

        loss = 0

        pos_loss = alpha * torch.log(pred) * torch.pow(1 - pred, 2) * pos_inds
        if weights is not None:
            pos_loss = pos_loss * weights[:-1]

        neg_loss = (1 - alpha) * torch.log(1 - pred) * torch.pow(pred, 2) * neg_inds

        num_pos = pos_inds.float().sum()
        pos_loss = pos_loss.sum()
        neg_loss = neg_loss.sum()

        if num_pos == 0:
            loss = loss - neg_loss
        else:
            loss = loss - (pos_loss + neg_loss) / num_pos
        return loss

    def _get_src_permutation_idx(self, indices):
        # permute predictions following indices
        batch_idx = torch.cat([torch.full_like(src, i) for i, (src, _) in enumerate(indices)])
        src_idx = torch.cat([src for (src, _) in indices])
        return batch_idx, src_idx

    def get_loss(self, loss, outputs, targets, indices, num, **kwargs):
        if 'pred_hoi_logits' in outputs.keys():
            loss_map = {
                'hoi_labels': self.loss_hoi_labels,
                'obj_labels': self.loss_obj_labels,
                'sub_obj_boxes': self.loss_sub_obj_boxes,
                'feats_mimic': self.mimic_loss
            }
        else:
            loss_map = {
                'obj_labels': self.loss_obj_labels,
                'obj_cardinality': self.loss_obj_cardinality,
                'verb_labels': self.loss_verb_labels,
                'sub_obj_boxes': self.loss_sub_obj_boxes
            }
        assert loss in loss_map, f'do you really want to compute {loss} loss?'
        return loss_map[loss](outputs, targets, indices, num, **kwargs)

    def forward(self, outputs, targets):
        outputs_without_aux = {k: v for k, v in outputs.items() if k != 'aux_outputs'}

        # Retrieve the matching between the outputs of the last layer and the targets
        indices = self.matcher(outputs_without_aux, targets)

        num_interactions = sum(len(t['hoi_labels']) for t in targets)
        num_interactions = torch.as_tensor([num_interactions], dtype=torch.float,
                                           device=next(iter(outputs.values())).device)
        if is_dist_avail_and_initialized():
            torch.distributed.all_reduce(num_interactions)
        num_interactions = torch.clamp(num_interactions / get_world_size(), min=1).item()

        # Compute all the requested losses
        losses = {}
        for loss in self.losses:
            losses.update(self.get_loss(loss, outputs, targets, indices, num_interactions))

        # In case of auxiliary losses, we repeat this process with the output of each intermediate layer.
        if 'aux_outputs' in outputs:
            for i, aux_outputs in enumerate(outputs['aux_outputs']):
                indices = self.matcher(aux_outputs, targets)
                for loss in self.losses:
                    kwargs = {}
                    if loss == 'obj_labels':
                        # Logging is enabled only for the last layer
                        kwargs = {'log': False}
                    l_dict = self.get_loss(loss, aux_outputs, targets, indices, num_interactions, **kwargs)
                    l_dict = {k + f'_{i}': v for k, v in l_dict.items()}
                    losses.update(l_dict)

        return losses

class PostProcessHOITriplet(nn.Module):

    def __init__(self, args):
        super().__init__()
        self.subject_category_id = args.subject_category_id

    @torch.no_grad()
    def forward(self, outputs, target_sizes):
        out_hoi_logits = outputs['pred_hoi_logits']
        out_obj_logits = outputs['pred_obj_logits']
        out_sub_boxes = outputs['pred_sub_boxes']
        out_obj_boxes = outputs['pred_obj_boxes']

        assert len(out_hoi_logits) == len(target_sizes)
        assert target_sizes.shape[1] == 2

        hoi_scores = out_hoi_logits.sigmoid()
        obj_scores = out_obj_logits.sigmoid()
        obj_labels = F.softmax(out_obj_logits, -1)[..., :-1].max(-1)[1]

        img_h, img_w = target_sizes.unbind(1)
        scale_fct = torch.stack([img_w, img_h, img_w, img_h], dim=1).to(hoi_scores.device)
        sub_boxes = box_cxcywh_to_xyxy(out_sub_boxes)
        sub_boxes = sub_boxes * scale_fct[:, None, :]
        obj_boxes = box_cxcywh_to_xyxy(out_obj_boxes)
        obj_boxes = obj_boxes * scale_fct[:, None, :]

        results = []
        for index in range(len(hoi_scores)):
            hs, os, ol, sb, ob = hoi_scores[index], obj_scores[index], obj_labels[index], sub_boxes[index], obj_boxes[
                index]
            sl = torch.full_like(ol, self.subject_category_id)
            l = torch.cat((sl, ol))
            b = torch.cat((sb, ob))
            results.append({'labels': l.to('cpu'), 'boxes': b.to('cpu')})

            ids = torch.arange(b.shape[0])

            results[-1].update({'hoi_scores': hs.to('cpu'), 'obj_scores': os.to('cpu'),
                                'sub_ids': ids[:ids.shape[0] // 2], 'obj_ids': ids[ids.shape[0] // 2:]})

        return results

def build_diffhoi_l(args):
    device = torch.device(args.device)

    num_classes = args.num_classes
    device = torch.device(args.device)

    backbone = build_backbone(args)

    transformer = build_deformable_transformer(args)

    try:
        match_unstable_error = args.match_unstable_error
        dn_labelbook_size = args.dn_labelbook_size
    except:
        match_unstable_error = True
        dn_labelbook_size = num_classes

    try:
        dec_pred_class_embed_share = args.dec_pred_class_embed_share
    except:
        dec_pred_class_embed_share = True
    try:
        dec_pred_bbox_embed_share = args.dec_pred_bbox_embed_share
    except:
        dec_pred_bbox_embed_share = True

    model = DiffHOI_L(
        backbone,
        transformer,
        num_classes=num_classes,
        num_queries=args.num_queries,
        aux_loss=True,
        iter_update=True,
        query_dim=4,
        random_refpoints_xy=args.random_refpoints_xy,
        fix_refpoints_hw=args.fix_refpoints_hw,
        num_feature_levels=args.num_feature_levels,
        nheads=args.nheads,
        dec_pred_class_embed_share=dec_pred_class_embed_share,
        dec_pred_bbox_embed_share=dec_pred_bbox_embed_share,
        # two stage
        two_stage_type=args.two_stage_type,
        # box_share
        two_stage_bbox_embed_share=args.two_stage_bbox_embed_share,
        two_stage_class_embed_share=args.two_stage_class_embed_share,
        decoder_sa_type=args.decoder_sa_type,
        num_patterns=args.num_patterns,
        dn_number = args.dn_number if args.use_dn else 0,
        dn_box_noise_scale = args.dn_box_noise_scale,
        dn_label_noise_ratio = args.dn_label_noise_ratio,
        dn_labelbook_size = dn_labelbook_size,
        args=args
    )
    matcher = build_matcher(args)
    weight_dict = {}
    if args.with_clip_label:
        weight_dict['loss_hoi_labels'] = args.hoi_loss_coef
        weight_dict['loss_obj_ce'] = args.obj_loss_coef
    else:
        weight_dict['loss_hoi_labels'] = args.hoi_loss_coef
        weight_dict['loss_obj_ce'] = args.obj_loss_coef

    weight_dict['loss_sub_bbox'] = args.bbox_loss_coef
    weight_dict['loss_obj_bbox'] = args.bbox_loss_coef
    weight_dict['loss_sub_giou'] = args.giou_loss_coef
    weight_dict['loss_obj_giou'] = args.giou_loss_coef
    # if args.with_mimic:
    #     weight_dict['loss_feat_mimic'] = args.mimic_loss_coef
    if args.aux_loss:
        aux_weight_dict = {}
        for i in range(args.dec_layers - 1):
            aux_weight_dict.update({k + f'_{i}': v for k, v in weight_dict.items()})
        weight_dict.update(aux_weight_dict)
    losses = ['hoi_labels', 'obj_labels', 'sub_obj_boxes']
    if args.with_mimic:
        losses.append('feats_mimic')

    criterion = SetCriterionHOI(args.num_obj_classes, args.num_queries, args.num_verb_classes, matcher=matcher,
                                weight_dict=weight_dict, eos_coef=args.eos_coef, losses=losses,
                                args=args)
    criterion.to(device)
    postprocessors = {'hoi': PostProcessHOITriplet(args)}

    return model, criterion, postprocessors


import networkx as nx
import matplotlib.pyplot as plt
import numpy as np
import itertools

def draw_graph(G: nx.Graph, path="Untitled"):
    pos = nx.spring_layout(G, k=1.5, scale=100)
    nx.draw_networkx_nodes(G, pos, node_color="blue")
    nx.draw_networkx_labels(G, pos, font_color="white")
    nx.draw_networkx_edges(G, pos, edge_color="gray")
    edge_labels = {(u, v): np.round(d["weight"], 2) for u, v, d in G.edges(data=True)}
    nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels, font_color="black", alpha=1,
                                 bbox=dict(boxstyle='round',
                                           ec=(1.0, 1.0, 1.0),
                                           fc=(1.0, 1.0, 1.0),
                                           alpha=0.0
                                           ))
    plt.savefig(path)

def get_nx_graph_by_adjacent_mat(_adjacent_mat, weight_amplifier=100, weight_cutoff=0.1):
    adjacent_mat = np.array(_adjacent_mat)
    adjacent_mat *= weight_amplifier
    G = nx.Graph(adjacent_mat)
    zero_edge = []
    for u, v, d in G.edges(data=True):
        if abs(d["weight"]) < weight_cutoff:
            zero_edge.append((u, v))
    for u, v in zero_edge:
        G.remove_edge(u, v)
    return G

def find_paths(g: nx.Graph):
    nodes = g.nodes
    combinations = itertools.combinations(nodes, 2)
    path_sets = set()
    paths = []
    for item in combinations:
        try:
            simple_paths = nx.algorithms.all_simple_paths(g, item[0], item[1])
            for path in simple_paths:
                path_sets.add(tuple(path))
                for i in range(len(path)):
                    sub_path = path[:i]
                    if len(sub_path) > 1:
                        path_sets.add(tuple(sub_path))
        except Exception as ex:
            pass
    for path in path_sets:
        paths.append(list(path))
    return paths


#!/bin/env python3

import re
import requests

first_url = "http://www.win4000.com/meinv197289_"
last_url = ""
for num in range(1,12):
    # 1:构造需要爬取图片的网址
    last_url = str(num)+".html"
    base_url = first_url + last_url

    # 2:解析出图片下载地址
    content = requests.get(base_url).text
    url = re.findall('img.*?class="pic-large".*?url="(.*?)".*?', content, re.S)
    url = url[0]

    # 3:获取图片并保存
    meitu = requests.get(url)

    with open(str(num)+'.jpg', 'wb') as f:
        f.write(meitu.content)
        f.close()


try:
    from setuptools import setup
except ImportError:
    from distutils.core import setup

config={
    'description':'season ticket breaker in Python',
    'author': 'Joe Zhou',
    'url': 'https://github.com/joezhouOxford/points.git',
    'download_url':'https://github.com/joezhouOxford/points.git',
    'author_email': 'joezhou.contact@gmail.com',
    'version': '0.1',
    'install_requires': ['nose','coverage'],
    'packages': ['points'],
    'scripts': [],
    'name': 'points'
}

setup(**config)


import pygame as pg
from settings import *
from sprites import *

class Game():
    def __init__(self):

        self.running = True
        pg.init()
        pg.mixer.init()
        self.screen = pg.display.set_mode((WIDTH, HEIGHT))
        pg.display.set_caption(title)
        self.clock = pg.time.Clock()
        self.font_name = pg.font.match_font(FONT_NAME)
        self.load_data()

    def load_data(self):
        pass

    def new(self):
        self.player1points = 0
        self.player2points = 0
        self.all_Sprites = pg.sprite.Group()
        self.players_group = pg.sprite.Group()
        self.ball_group = pg.sprite.Group()

        self.player1 = Player1(self)
        self.player2 = Player2(self)
        self.ball = Ball(self)

        self.all_Sprites.add(self.player1)
        self.players_group.add(self.player1)
        self.all_Sprites.add(self.player2)
        self.players_group.add(self.player2)
        for i in range(1):
            self.all_Sprites.add(self.ball)
            self.ball_group.add(self.ball)

    def run(self):
        self.playing = True
        while self.playing:
            self.dt = self.clock.tick(FPS)
            self.event()
            self.update()
            self.draw()

    def event(self):
        for event in pg.event.get():
            if event.type == pg.QUIT:
                if self.playing:
                    self.playing = False
                self.running = False

    def update(self):
        self.all_Sprites.update()

        hits = pg.sprite.collide_rect(self.ball,self.player1)
        if hits:
            self.ball.speedx *= -1
            self.ball.speedy *= 1
        hits = pg.sprite.collide_rect(self.ball, self.player2)
        if hits:
            self.ball.speedx *= -1
            self.ball.speedy *= 1

    def spawn_ball(self):
        for i in range(1):
            self.ball_group.add(self.ball)
            self.all_Sprites.add(self.ball)
        self.wait_for_key()

    def draw_text(self,text,size,color,x,y):
        font = pg.font.Font(self.font_name,size)
        text_surface = font.render(text,True,color)
        text_rect = text_surface.get_rect()
        text_rect.midtop = (x,y)
        self.screen.blit(text_surface,text_rect)

    def draw(self):
        self.screen.fill(BLACK)
        self.draw_text("Score:" +str(self.player1points),25,WHITE,WIDTH-50,15)
        self.draw_text("Score:" + str(self.player2points), 25, WHITE,50, 15)
        self.draw_text("PONG", 25, WHITE, WIDTH/2, 15)
        self.all_Sprites.draw(self.screen)

        pg.display.flip()

    def wait_for_key(self):
        waiting = True
        while waiting:
            self.clock.tick()
            for event in pg.event.get():
                if event.type == pg.QUIT:
                    pg.quit()
                if event.type == pg.KEYUP:
                    waiting = False
                    break

g = Game()
while g.running:
    g.wait_for_key()
    g.new()
    g.run()

pg.quit()


class Currency:
    """This is the currency class. It stores the currency name, currency code
    and country name as attributes"""

    def __init__(self, currencyName, currencyCode, countryName):
        self.currencyName=currencyName
        self.currencyCode=currencyCode
        self.countryName=countryName


# DESCRIPTION: This file implemented a biased l2-norm regularized SVD model with SGD. A bias ci is added for each user i and dj for each item j. The prediction function is U[i]*V[j].T+c[i]+d[j] for user i and item j. SGD terminates when the validation error stops decreasing.

# USAGE: To train the model, run "python3 code/RegularizedSVD2.py -k=5" and "python3 code/RegularizedSVD2.py -k=5 -d=1". "-k" specifies the number of dimensions for dimension reduction in SVD. "-d" chooses the training/validation data split.

import Initialization
import SVD
import numpy as np
import Globals
import random
import time
from sklearn import linear_model

def biasedRSVD(train,test,k=96):
	# initialization
	print('start initialization k =',k)
	lrate = Globals.lrate
	lamb = 0.02
	lamb2 = 0.05
	mu = 0
	sigma = 1
	suffix = '_fixed'+Globals.dataIdx+'.npy'
	if not Globals.fixed:
		suffix = '.npy'
	# read in previously trained result
	if Globals.warmStart:
		print('warm start')
		U = np.load('./log/RSVD2_U_'+str(k)+Globals.modelIdx+suffix)
		Vt = np.load('./log/RSVD2_Vt_'+str(k)+Globals.modelIdx+suffix)
		c = np.load('./log/RSVD2_c_'+str(k)+Globals.modelIdx+suffix)
		d = np.load('./log/RSVD2_d_'+str(k)+Globals.modelIdx+suffix)
	# otherwise, random initialization
	else:
		U = np.random.rand(Globals.nUsers,k)
		Vt = np.random.rand(k,Globals.nItems)
		c = np.random.rand(Globals.nUsers)
		d = np.random.rand(Globals.nItems)
	known = train!=0
	print('finish initialization')

	print('start SGD')
	startTime = time.time()
	globalMean = np.mean(train[known])
	t = 0
	prev1 = 1000000
	prev2 = 1000000
	while True:
		# random choice of training sample
		i = random.randint(0,Globals.nUsers-1)
		j = random.randint(0,Globals.nItems-1)
		while not known[i,j]:
			i = random.randint(0,Globals.nUsers-1)
			j = random.randint(0,Globals.nItems-1)

		# gradient descent
		yp = c[i]+d[j]+U[i,:].dot(Vt[:,j])
		r = train[i,j] - yp
		Ut = U[i,:].T
		U[i,:] += lrate*(r*Vt[:,j].T-lamb*U[i,:])
		Vt[:,j] += lrate*(r*Ut-lamb*Vt[:,j])
		tmp = lrate*(r - lamb2*(c[i]+d[j]-globalMean))
		c[i] += tmp
		d[j] += tmp

		# evaluation
		if t%10000 == 0:
			A = U.dot(Vt)
			C = np.reshape(c,(Globals.nUsers,1))
			D = np.reshape(d,(1,Globals.nItems))
			C = np.repeat(C,Globals.nItems,axis=1)
			D = np.repeat(D,Globals.nUsers,axis=0)
			A += C+D
			score = SVD.evaluation2(A,test)
			print('t =',t,'score =',score)
			if score > prev2 and prev2 > prev1:
				if lrate <= 1e-5:
					break
				else:
					lrate *= 0.1
					print('learning rate =',lrate)
			prev1 = prev2
			prev2 = score

		# auto save
		if t%500000 == 0:
			np.save('./log/RSVD2_U_'+str(k)+Globals.modelIdx+suffix,U)
			np.save('./log/RSVD2_Vt_'+str(k)+Globals.modelIdx+suffix,Vt)
			np.save('./log/RSVD2_c_'+str(k)+Globals.modelIdx+suffix,c)
			np.save('./log/RSVD2_d_'+str(k)+Globals.modelIdx+suffix,d)
			print('intermediate result saved')
		t += 1
	endTime = time.time()
	print('finish SGD',int(endTime-startTime),'s')
	np.save('./log/RSVD2_U_'+str(k)+Globals.modelIdx+suffix,U)
	np.save('./log/RSVD2_Vt_'+str(k)+Globals.modelIdx+suffix,Vt)
	np.save('./log/RSVD2_c_'+str(k)+Globals.modelIdx+suffix,c)
	np.save('./log/RSVD2_d_'+str(k)+Globals.modelIdx+suffix,d)

	# clipping
	print('start clipping')
	A = U.dot(Vt)
	C = np.reshape(c,(Globals.nUsers,1))
	D = np.reshape(d,(1,Globals.nItems))
	C = np.repeat(C,Globals.nItems,axis=1)
	D = np.repeat(D,Globals.nUsers,axis=0)
	A += C+D
	# over 5
	mask = A>5
	A[mask] = 5
	# below 1
	mask = A<1
	A[mask] = 1
	print('finish clipping')
	score = SVD.evaluation2(A,test)
	print('after clipping score =',score)
	return A

# choose the best number of dimensions kept
def chooseK(train,test):
	for k in range(25,4,-5):
		biasedRSVD(train,test,k)

if __name__ == "__main__":
	Initialization.initialization()
	train, test = Initialization.readInData2(idx=Globals.dataIdx)
	# choose the best k
	if Globals.predict == 'k':
		chooseK(train,test)
	# train & predict
	else:
		A = biasedRSVD(train,test,Globals.k)
		np.save('./log/RSVD2_A_'+str(Globals.k)+'_fixed'+Globals.dataIdx+'.npy',A)