import pygame
import config
import math
from unit import Unit
from unit_move import UnitMove

class Enemy(Unit):
    def __init__(self, x, y):
        super().__init__(x, y)
        self.start_position = [x, y]
        self.time_till_damage = 0
        self.look = []              # [center_x, center_y, radius]
        self.score = 0

    def move_to_position(self, move_to_position):
        position = self.get_position()

        # Find direction vector (dx, dy) between enemy and player.
        dx, dy = move_to_position[0] - position[0], move_to_position[1] - position[1]
        dist = math.hypot(dx, dy)

        if dx < 0:
            self.set_direction(UnitMove.LEFT)
        else:
            self.set_direction(UnitMove.RIGHT)

        if dist > 1:
            dx, dy = dx / dist, dy / dist  # Normalize.
            # Move along this normalized vector towards the player at current speed.
            position[0] += dx * config.BAT_VELOCITY
            position[1] += dy * config.BAT_VELOCITY
            self.set_position(position[0], position[1])

    def contains_look(self, player):
        player_corners = player.get_hitbox_corners()
        enemy_look = self.get_look()
        player_center = [player_corners[0][0] + (player_corners[1][0] - player_corners[0][0])/2,
                         player_corners[0][1] + (player_corners[3][1] - player_corners[0][1])/2]
        dx = enemy_look[0] - player_center[0]
        dy = enemy_look[1] - player_center[1]
        if dx * dx + dy * dy <= enemy_look[2] * enemy_look[2]:
            return True
        return False

    def contains(self, player):
        player_hitbox = player.get_hitbox()
        enemy_hitbox = self.get_hitbox()
        time_till_damage = self.get_time_till_damage()
        if time_till_damage == 0:
            # Check player up left corner is in enemy hitbox
            if enemy_hitbox[0] <= player_hitbox[0] <= enemy_hitbox[0] + enemy_hitbox[2]:
                if enemy_hitbox[1] <= player_hitbox[1] <= enemy_hitbox[1] + enemy_hitbox[3]:
                    self.set_time_till_damage(time_till_damage + 1)
                    return True

            # Check player up right corner is in enemy hitbox
            if enemy_hitbox[0] <= player_hitbox[0] <= enemy_hitbox[0] + enemy_hitbox[2]:
                if enemy_hitbox[1] <= player_hitbox[1] + player_hitbox[3] <= enemy_hitbox[1] + enemy_hitbox[3]:
                    self.set_time_till_damage(time_till_damage + 1)
                    return True

            # Check player down right corner is in enemy hitbox
            if enemy_hitbox[0] <= player_hitbox[0] + player_hitbox[2] <= enemy_hitbox[0] + enemy_hitbox[2]:
                if enemy_hitbox[1] <= player_hitbox[1] + player_hitbox[3] <= enemy_hitbox[1] + enemy_hitbox[3]:
                    self.set_time_till_damage(time_till_damage + 1)
                    return True

            # Check player down left corner is in enemy hitbox
            if enemy_hitbox[0] <= player_hitbox[0] + player_hitbox[2] <= enemy_hitbox[0] + enemy_hitbox[2]:
                if enemy_hitbox[1] <= player_hitbox[1] <= enemy_hitbox[1] + enemy_hitbox[3]:
                    self.set_time_till_damage(time_till_damage + 1)
                    return True
        else:
            time_till_damage += 1
            if time_till_damage < config.TIME_TILL_DAMAGE:
                self.set_time_till_damage(time_till_damage + 1)
            else:
                self.set_time_till_damage(0)

        return False

    def render_look(self, screen, camera):
        look = self.get_look()
        pygame.draw.circle(screen, (0, 0, 255), [look[0] - camera[0], look[1] - camera[1]], look[2], 1)

    def set_look(self, hitbox):
        self.look = [hitbox[0] + hitbox[2] / 2, hitbox[1] + hitbox[3] / 2, config.RADIUS]

    def get_look(self):
        return self.look

    def get_start_position(self):
        return self.start_position

    def set_time_till_damage(self, time_till_damage):
        self.time_till_damage = time_till_damage

    def get_time_till_damage(self):
        return self.time_till_damage

    def move_directrion(self, dx, dy):
        pass

    def get_score(self):
        return self.score


import typing as t

import numpy as np
from emo_utils import convert_to_one_hot
from emo_utils import predict
from emo_utils import softmax
from tensorflow.keras.layers import LSTM
from tensorflow.keras.layers import Activation
from tensorflow.keras.layers import Dense
from tensorflow.keras.layers import Dropout
from tensorflow.keras.layers import Embedding
from tensorflow.keras.layers import Input
from tensorflow.keras.models import Model

def sentence_to_avg(
    sentence: str,
    word_to_vec_map: dict[str, t.Any],
) -> np.ndarray:
    """
    Converts a sentence (string) into a list of words (strings). Extracts the GloVe
    representation of each word and averages its value into a single vector encoding
    the meaning of the sentence.

    Arguments:
    sentence -- string, one training example from X
    word_to_vec_map -- dictionary mapping every word in a vocabulary into its
        50-dimensional vector representation

    Returns:
    avg -- average vector encoding information about the sentence, numpy-array of
        shape (J,), where J can be any number
    """

    words = [w.lower() for w in sentence.split()]

    any_word = list(word_to_vec_map.keys())[0]
    avg = np.zeros(word_to_vec_map[any_word].shape[0])

    count = 0

    for w in words:
        if w in word_to_vec_map:
            avg += word_to_vec_map[w]
            count += 1

    if count > 0:
        avg = avg / count

    return avg

def model(
    X,
    Y,
    word_to_vec_map,
    learning_rate=0.01,
    num_iterations=400,
):
    """
    Model to train word vector representations in numpy.

    Arguments:
        X -- input data, numpy array of sentences as strings, of shape (m,)
        Y -- labels, numpy array of integers between 0 and 7, numpy-array of shape (m, 1)
        word_to_vec_map -- dictionary mapping every word in a vocabulary into its
            50-dimensional vector representation
        learning_rate -- learning_rate for the stochastic gradient descent algorithm
        num_iterations -- number of iterations

    Returns:
        pred -- vector of predictions, numpy-array of shape (m, 1)
        W -- weight matrix of the softmax layer, of shape (n_y, n_h)
        b -- bias of the softmax layer, of shape (n_y,)
    """

    # Get a valid word contained in the word_to_vec_map
    any_word = list(word_to_vec_map.keys())[0]

    # number of training examples
    m = Y.shape[0]
    # number of classes
    n_y = len(np.unique(Y))
    # dimensions of the GloVe vectors
    n_h = word_to_vec_map[any_word].shape[0]

    # Initialize parameters using Xavier initialization
    W = np.random.randn(n_y, n_h) / np.sqrt(n_h)
    b = np.zeros((n_y,))

    # Convert Y to Y_one_hot with n_y classes
    Y_oh = convert_to_one_hot(Y, C=n_y)

    # Optimization loop
    for t in range(num_iterations):
        cost = 0
        dW = 0
        db = 0

        # Loop over the training examples
        for i in range(m):
            # Average the word vectors of the words from the i'th training example
            avg = sentence_to_avg(X[i], word_to_vec_map)

            # Forward propagate the avg through the softmax layer.
            z = W @ avg + b
            a = softmax(z)

            # Add the cost using the i'th training label's one hot representation and
            # "A" (the output of the softmax)
            cost += -np.sum(Y_oh[i] * np.log(a))

            # Compute gradients
            dz = a - Y_oh[i]
            dW += np.dot(dz.reshape(n_y, 1), avg.reshape(1, n_h))
            db += dz

            # Update parameters with Stochastic Gradient Descent
            W = W - learning_rate * dW
            b = b - learning_rate * db

        assert type(cost) == np.float64, "Incorrect implementation of cost"
        assert cost.shape == (), "Incorrect implementation of cost"

        if t % 100 == 0:
            print("Epoch: " + str(t) + " --- cost = " + str(cost))
            pred = predict(X, Y, W, b, word_to_vec_map)

    return pred, W, b

def sentences_to_indices(
    X,
    word_to_index,
    max_len,
):
    """
    Converts an array of sentences (strings) into an array of indices corresponding to
    words in the sentences. The output shape should be such that it can be given to
    `Embedding()` (described in Figure 4).

    Arguments:
    X -- array of sentences (strings), of shape (m,)
    word_to_index -- a dictionary containing the each word mapped to its index
    max_len -- maximum number of words in a sentence. You can assume every sentence in
        X is no longer than this.

    Returns:
    X_indices -- array of indices corresponding to words in the sentences from X, of
        shape (m, max_len)
    """

    # number of training examples
    m = X.shape[0]

    X_indices = np.zeros((m, max_len))

    for i in range(m):
        sentence_words = [w.lower() for w in X[i].split()]

        j = 0
        for w in sentence_words:
            if w in word_to_index:
                X_indices[i, j] = word_to_index[w]
                j += 1

    return X_indices

def pretrained_embedding_layer(
    word_to_vec_map,
    word_to_index,
):
    """
    Creates a Keras Embedding() layer and loads in pre-trained GloVe 50-dimensional
    vectors.

    Arguments:
        word_to_vec_map -- dictionary mapping words to their GloVe vector representation.
        word_to_index -- dictionary mapping from words to their indices in
            the vocabulary (400,001 words)

    Returns:
        embedding_layer -- pretrained layer Keras instance
    """
    # adding 1 to fit Keras embedding (requirement)
    vocab_size = len(word_to_index) + 1
    any_word = list(word_to_vec_map.keys())[0]
    # define dimensionality of your GloVe word vectors (= 50)
    emb_dim = word_to_vec_map[any_word].shape[0]

    # Initialize the embedding matrix as a numpy array of zeros.
    emb_matrix = np.zeros((vocab_size, emb_dim))

    # Set each row "idx" of the embedding matrix to be
    # the word vector representation of the idx'th word of the vocabulary
    for word, idx in word_to_index.items():
        emb_matrix[idx, :] = word_to_vec_map[word]

    # Define Keras embedding layer with the correct input and output sizes
    # Make it non-trainable.
    embedding_layer = Embedding(
        input_dim=vocab_size,
        output_dim=emb_dim,
        trainable=False,
    )
    embedding_layer.build((None,))

    # Set the weights of the embedding layer to the embedding matrix.
    embedding_layer.set_weights([emb_matrix])

    return embedding_layer

def Emojify_V2(
    input_shape,
    word_to_vec_map,
    word_to_index,
):
    """
    Function creating the Emojify-v2 model's graph.

    Arguments:
        input_shape -- shape of the input, usually (max_len,)
        word_to_vec_map -- dictionary mapping every word in a vocabulary into its
            50-dimensional vector representation
        word_to_index -- dictionary mapping from words to their indices in the
            vocabulary (400,001 words)

    Returns:
        model -- a model instance in Keras
    """

    ### START CODE HERE ###
    # Define sentence_indices as the input of the graph.
    # It should be of shape input_shape and dtype 'int32'
    # (as it contains indices, which are integers).
    sentence_indices = Input(shape=input_shape, dtype="int32")

    # Create the embedding layer pretrained with GloVe Vectors (≈1 line)
    embedding_layer = pretrained_embedding_layer(word_to_vec_map, word_to_index)

    # Propagate sentence_indices through your embedding layer
    # (See additional hints in the instructions).
    embeddings = embedding_layer(sentence_indices)

    # Propagate the embeddings through an LSTM layer with 128-dimensional hidden state
    # The returned output should be a batch of sequences.
    X = LSTM(units=128, return_sequences=True)(embeddings)
    # Add dropout with a probability of 0.5
    X = Dropout(rate=0.5)(X)
    # Propagate X trough another LSTM layer with 128-dimensional hidden state
    # The returned output should be a single hidden state, not a batch of sequences.
    X = LSTM(units=128, return_sequences=False)(X)
    # Add dropout with a probability of 0.5
    X = Dropout(rate=0.5)(X)
    # Propagate X through a Dense layer with 5 units
    X = Dense(units=5)(X)
    # Add a softmax activation
    X = Activation("softmax")(X)

    # Create Model instance which converts sentence_indices into X.
    model = Model(inputs=sentence_indices, outputs=X)

    return model

Course 6, Sequence Models/Week 2/assignment_2/Emoji_v3a.py


import subprocess
import threading
import io
from fcntl import fcntl, F_GETFL, F_SETFL
from os import O_NONBLOCK
import sys
#from flask_socketio import SocketIO

command = "pintos -v -k --qemu --disk cs162proj.dsk -- -q run shell"

class Shell():
    def set_flags(self, pipe):
        flags = fcntl(pipe, F_GETFL)
        fcntl(pipe, F_SETFL, flags | O_NONBLOCK)

    def __init__(self, app, command):
        self.app = app
        self.cmd = command
        self.p = None
        self.output_thread = threading.Thread()

    def run(self):
        if self.p == None:
            self.p = subprocess.Popen(self.cmd.split(' '),
                    cwd="./os_build",
                    stdin=subprocess.PIPE,
                    stdout=subprocess.PIPE,
                    )
            self.set_flags(self.p.stdout)
            self.set_flags(self.p.stdin)

    def output(self):
        buf = io.StringIO()
        while True:
            data = self.p.stdout.read(1)
            if data:
                buf.write(data.decode('utf-8'))
            else:
                if (buf.getvalue() != ""):
                    self.app.emit('send_output', buf.getvalue())
                    buf = io.StringIO()

    def input(self, command):
        self.p.stdin.write(command.encode())
        self.p.stdin.flush()

shell = Shell(None, command)


#!/usr/bin/env python3
#
# 10. Bayesian History Matching technique (advanced use)
#
import os
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import json

from GPErks.constants import DEFAULT_TMP_OUTFILE_DIR
from GPErks.perks.history_matching import Wave
from GPErks.serialization.labels import read_labels_from_file
from GPErks.serialization.path import posix_path
from GPErks.utils.array import get_minmax
from GPErks.utils.plotting import interp_col, get_col
from GPErks.utils.sampling import Sampler
from gpytorch.kernels import MaternKernel, ScaleKernel
from gpytorch.likelihoods import GaussianLikelihood
from torchmetrics import MeanSquaredError, R2Score

from GPErks.gp.data.dataset import Dataset
from GPErks.gp.experiment import GPExperiment
from GPErks.gp.mean import LinearMean
from GPErks.log.logger import get_logger
from GPErks.train.emulator import GPEmulator
from GPErks.utils.random import set_seed

def main():
    # Set logger and enforce reproducibility
    log = get_logger()
    seed = 8
    set_seed(seed)

    # Load experimental values (mean and variance) you aim to match
    exp_data_file = posix_path(os.getcwd(), "data", "example_10", "expdata.json")
    expdata = {}
    with open(exp_data_file, "r") as f:
        expdata = json.load(f)
    exp_mean = [val["mean"] for val in expdata.values()]
    exp_var = [val["var"] for val in expdata.values()]

    # Load input parameters and output features' names
    dataset_dir = Path(posix_path(os.getcwd(), "datasets", "stefano", "8p", "sham"))
    xlabels = read_labels_from_file(dataset_dir / "xlabels.txt")
    ylabels = read_labels_from_file(dataset_dir / "ylabels.txt")
    feature_idx = {key: val for val, key in enumerate(ylabels)}
    active_features = list(expdata.keys())
    active_indices = [feature_idx[f] for f in active_features]

    # Train list of univariate emulators (one for each feature to match)
    X = np.loadtxt(dataset_dir / "X.txt", dtype=float)
    Y = np.loadtxt(dataset_dir / "Y.txt", dtype=float)

    emulators = []
    for idx, feature in zip(active_indices, active_features):
        y = Y[:, idx]

        dataset = Dataset(X, y, x_labels=xlabels, y_label=feature)
        likelihood = GaussianLikelihood()
        mean = LinearMean(degree=1, input_size=dataset.input_size, bias=True)
        covar = ScaleKernel(MaternKernel(ard_num_dims=dataset.input_size))
        metrics = [MeanSquaredError(), R2Score()]

        experiment = GPExperiment(
            dataset,
            likelihood,
            mean,
            covar,
            metrics=metrics,
            seed=seed
        )
        device = "cpu"

        emulator = GPEmulator(experiment, device)
        emulator.train_auto()
        emulators.append(emulator)

    minmax = get_minmax(X)

    waveno = 1  # number of iteration we are at (wave id if you want)
    cutoff = 3.0  # threshold value for the implausibility criterion
    maxno = 1  # explained below
    # The univariate GPE of each output feature will give for each point x_i a specific implausibility measure.
    # With the current implausibility criterion, for each x_i we take the maximum implausibility across all the output
    # features. With maxno=1, the maximum is calculated across all the output features (i.e., till the last worse
    # implausibility measure). If maxno=2 --> till the previous-to-last worse implausibility measure and so on.
    # With this criterion, the worse-performing emulator (the output feature which is the least well captured) will
    # dominate the entire analysis and thus determine if a point is non-implausible or implausible

    w = Wave(
        emulator=emulators,
        Itrain=minmax,
        cutoff=cutoff,
        maxno=maxno,
        mean=exp_mean,
        var=exp_var,
    )  # instantiate the wave object

    sampler = Sampler(design="lhs", dim=X.shape[1], seed=seed)
    n_samples = 100000
    X = sampler.sample(
        n_samples,
        l_bounds=list(minmax[:, 0]),
        u_bounds=list(minmax[:, 1]),
    )

    # Run one iteration of HM, which is: apply the implausibility criterion to detect regions of non-implausible
    # and of implausible points starting from the initial samples in X
    w.find_regions(X)
    w.print_stats()  # show statistics about the two obtained spaces
    w.plot_wave(xlabels=xlabels, display="impl")  # plot the current wave of history matching (impl. measure plot)
    w.plot_wave(xlabels=xlabels, display="var")  # we can also check the accuracy of the GPEs for the current wave
    # note: if filepath=<path_to_file> flag is provided, the plot will be saved to <path_to_file>

    # How to continue on the next wave in 5 steps
    #
    # (0) Save an exact copy of the wave. We always recommend saving each wave right on completion before manipulating
    # its internal structure as you might need it later for other purposes (see Appendix - A2)
    outfiles_dir = Path(DEFAULT_TMP_OUTFILE_DIR)
    outfiles_dir.mkdir(parents=True, exist_ok=True)

    w0 = w.copy()
    w0.print_stats()
    w0.save(outfiles_dir / f"wave_{waveno}.json")

    # (1) From the current non-implausible region, select points to be simulated and points to be used as tests
    # for the next wave
    n_tests = 10000  # number of TEST points we want for the next wave
    n_simuls = 128  # number of current NIMP points we want to simulate to augment training dataset for the next wave
    n_avail_nimps = len(w0.nimp_idx)  # we currently have available only this number of NIMP points
    if n_tests + n_simuls > n_avail_nimps:  # if they are not enough
        n_total_points = n_tests + n_simuls
        w.augment_nimp(n_total_points)  # use 'cloud technique' to generate new NIMP points starting from existing ones

    # Get the requested datasets
    X_simul, X_test = w.get_nimps(n_simuls)

    # We now have all the necessary data to run the next wave: a dataset to simulate to augment the training dataset
    # and build new emulators, and new TEST points to be evaluated with the new emulators. Saving the data to files.
    np.savetxt(outfiles_dir / f"X_simul_{waveno}.txt", X_simul, fmt="%.6f")
    np.savetxt(outfiles_dir / f"X_test_{waveno}.txt", X_test, fmt="%.6f")

    w.print_stats()  # quick check on TESTS, IMP, and NIMP sets' sizes after augmentation

    # (2) Simulate the selected points X_simul
    # (3) Add the simulated points and respective results to the training dataset used in the previous wave
    # (3) Train GPEs on the new, augmented training dataset
    # (4) Start a new wave of HM, where the initial parameter space to be split into non-implausible and
    # implausible regions is no more a Latin Hypercube design but is now the non-implausible region obtained
    # (and augmented) in the previous wave (i.e., X_test)

    # Appendix
    #
    # (A1) Visual check on the datasets generated for the next wave
    X_nimp = w.NIMP
    X_test = np.loadtxt(outfiles_dir / f"X_test_{waveno}.txt", dtype=float)
    X_simul = np.loadtxt(outfiles_dir / f"X_simul_{waveno}.txt", dtype=float)

    # We will inspect only 2 dimensions of the full 8D parameter space to keep it simple
    param = [4, 5]  # select 2 dimensions
    subset_idx = list(np.random.randint(0, X_test.shape[0], size=10*X_simul.shape[0]))  # select an example portion

    colors = interp_col(get_col("blue"), 4)  # getting some blue colour variants

    # Plotting current wave NIMP + next wave TEST + next wave SIMUL
    fig, axis = plt.subplots(1, 1)
    axis.scatter(X_nimp[:, param[0]], X_nimp[:, param[1]], fc=colors[1], ec=colors[1], label=f"X_nimp of wave {waveno}")
    axis.scatter(X_test[subset_idx, param[0]], X_test[subset_idx, param[1]], fc=colors[-1], ec=colors[-1], label=f"X_test for wave {waveno+1}")
    axis.scatter(X_simul[:, param[0]], X_simul[:, param[1]], fc='r', ec='r', label=f"X_simul for wave {waveno+1}")
    axis.set_xlabel(xlabels[param[0]], fontsize=12)
    axis.set_ylabel(xlabels[param[1]], fontsize=12)
    axis.legend()
    fig.tight_layout()
    plt.show()  # TEST + SIMUL points for NEXT wave are all within NIMP space CURRENT wave

    # (A2) Loading a wave object
    # You can load a wave object by providing the same data used to instantiate the wave (emulator, Itrain, cutoff,
    # maxno, mean, var). This is normally done when you need to re-run the wave differently. Alternatively, you can load
    # the wave object with no arguments. This is normally done when you need to examine the wave internal structure.
    # Let's try loading with no arguments.
    w = Wave()
    w.load(Path(outfiles_dir) / f"wave_{waveno}.json")
    w.print_stats()  # notice that TESTS, IMP, and NIMP sets' sizes are the same as pre-augmentation

    # You can get a list of all wave object attributes by printing:
    # print(w.__dict__.keys())

    # Noteworthy attributes are:
    # W.I: implausibility measure obtained for each point in the test set
    # W.PV: percentage emulator variance over experimental variance at each point (given as a fraction)
    # W.NIMP: non-implausible region
    # W.nimp_idx: indices of the initial test set which resulted to be non-implausible
    # W.IMP: implausible region
    # W.imp_idx: indices of the initial test set which resulted to be implausible
    # W.simul_idx: indices of W.NIMP that were selected to be simulated for the next wave
    # W.nsimul_idx: indices of W.NIMP which were not selected for simulations
    # (the respective points will appear in the test set of the next wave instead)

    # The original test set is not stored as an attribute to save space. However, this information can still be
    # retrieved from the stored attributes as:
    X_test = w.reconstruct_tests()
    print((np.equal(X_test, X)).all())  # the test set of first wave was the LHD we generated initially in this script

if __name__ == "__main__":
    main()


import sys
sys.path.append('../../preprocess')
from make_pca import load_landmarks
import numpy as np
import tensorflow as tf
from pfld import predict_landmarks as pfld_predict_landmarks
from pfld_custom import predict_landmarks as pfld_custom_predict_landmarks
from skimage.color import rgb2gray
import cv2
import dlib
from skimage.transform import resize
from prepare_data import IMAGE_SIZE, view_img, resize_lmks
from skimage.transform import resize
import matplotlib
from train_pfld import normalize_data
import os
from detector import get_face_detector

matplotlib.use("TkAgg")
# IMAGE_SIZE = 224

class Rect:
    def __init__(self, t, b, l, r):
        self.t = t
        self.b = b
        self.l = l
        self.r = r
    def top(self):
        return self.t
    def bottom(self):
        return self.b
    def right(self):
        return self.r
    def left(self):
        return self.l

def predict(data, model_path, predict_fn, image_size=IMAGE_SIZE, depth_multiplier=1.0, **kwargs):
    input_shape = [None, image_size, image_size, 3]
    inputs = tf.placeholder(tf.float32, shape=input_shape, name='input_images')
    preds, _, _ = predict_fn(inputs, image_size, is_training=False, depth_multiplier=depth_multiplier, **kwargs)
    print('predict tensor = ', preds)
    saver = tf.train.Saver()
    # g = tf.get_default_graph()
    # tf.contrib.quantize.create_eval_graph(input_graph=g)
    with tf.Session() as sess:
        saver.restore(sess, model_path)
        # sess.run(tf.global_variables_initializer())
        results = sess.run(preds, feed_dict={inputs: data})
        print('landmarks = ', results)
        # print('S1 > ')
        return results

def predict_tflite(data, model_path):
    interpreter = tf.lite.Interpreter(model_path=model_path)
    interpreter.allocate_tensors()
    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()
    print('input_details ', input_details[0], ' data shape ', data.shape)
    interpreter.set_tensor(input_details[0]['index'], data)
    interpreter.invoke()
    landmarks = interpreter.get_tensor(output_details[0]['index'])
    print('landmarks = ', landmarks)
    return landmarks

def crop(img, box):
    return img[box.top(): box.bottom(), box.left(): box.right()]

def crop_landmarks(landmarks, box):
    return landmarks - np.array([box.left(), box.top()])

def predict_single(img_path, model_path,
                image_size=IMAGE_SIZE,
                depth_multiplier=1.0,
                predict_fn=pfld_predict_landmarks,
                zero_mean=True,
                box_detector='dlib',
                **kwargs):
    img_size = image_size
    gt_landmark = None
    if box_detector == 'gt':
        points, imgs_sizes, imgs = load_landmarks('/home/tamvm/Downloads/ibug_300W_large_face_landmark_dataset/labels_ibug_300W_train.xml')
        fn = os.path.basename(img_path)
        gt_landmark = None
        for idx, img in enumerate(imgs):
            if img.endswith(fn):
                gt_landmark = points[idx]
                break
        if gt_landmark is not None:
            min_y, max_y = gt_landmark[:,1].min(), gt_landmark[:,1].max()
            min_x, max_x = gt_landmark[:,0].min(), gt_landmark[:,0].max()
            box = Rect(min_y, max_y, min_x, max_x)
            # _, gt_landmark = crop_and_resize(, gt_landmark, image_size)
    elif box_detector == 'tf':
        detector = get_face_detector()
        l, t, r, b = detector.detect(img_path)
        box = Rect(t, b, l, r)

    # get face bound
    else:
        img = dlib.load_rgb_image(img_path)
        detector = dlib.get_frontal_face_detector()
        box = detector(img, 1)[0]

    oridata = cv2.imread(img_path)
    # if image_size ==80:
    #     oridata = cv2.cvtColor(oridata,cv2.COLOR_BGR2RGB)
    data = crop(oridata, box)
    data = resize(data, (img_size, img_size), anti_aliasing=True, mode='reflect')
    # view_img(data, None)
    normalized_data = normalize_data(data)
    if model_path.endswith('.tflite'):
        # print('using tflite model ', model_path)
        # is_unint8 = model_path.find('uint8') >= 0
        # if is_unint8:
        #     print('int model')
        #     lmks = predict_tflite((np.reshape(data, (1, *data.shape)) * 255).astype(np.uint8), model_path)[0]
        # else:
        print('float model')
        lmks = predict_tflite(np.reshape(normalized_data, (1, *normalized_data.shape)).astype(np.float32), model_path)[0]
    else:
        lmks = predict(np.reshape(normalized_data, (1, *normalized_data.shape)), model_path, predict_fn,
                        image_size=image_size,
                        depth_multiplier=depth_multiplier,
                        **kwargs)[0]
    # print('landmark = ', lmks)
    if zero_mean:
        for i in range(0, 68):
            lmks[i*2] = (lmks[i*2]/2+0.5)* image_size# (lmks[i*2]/2+0.5)*image_size
            lmks[i*2+1] = (lmks[i*2+1]/2 + 0.5) * image_size# (lmks[i*2+1]/2 + 0.5)*image_size
    else:
        for i in range(0, 68):
            lmks[i*2] = (lmks[i*2])* image_size# (lmks[i*2]/2+0.5)*image_size
            lmks[i*2+1] = (lmks[i*2+1]) * image_size# (lmks[i*2+1]/2 + 0.5)*image_size
        # print('landmarks after denorm', lmks)
    lmks = lmks.reshape((68, 2))

    view_img(data, lmks)

if __name__ == '__main__':
    # 2960256451_1.jpg
    # '/home/tamvm/Downloads/ibug_300W_large_face_landmark_dataset/helen/testset/30427236_1.jpg'
    use_tflite = False
    model = 'pfld-custom-80-025m-saux7-x3'
    # model = 'ailab'
    if model == 'pfld-64':
        predict_single('/home/tamvm/Downloads/test_face_tamvm_2.jpg', #'/home/tamvm/Downloads/ibug_300W_large_face_landmark_dataset/helen/trainset/2960256451_1.jpg',
            '../../data/checkpoints-pfld-64-05m/pfld-311400' if not use_tflite else '../../data/pfld-64-quant.tflite',
            depth_multiplier=0.5,
            image_size=64)
    elif model == 'pfld-112':
        predict_single('/home/tamvm/Downloads/test_face_tamvm_2.jpg', #'/home/tamvm/Downloads/ibug_300W_large_face_landmark_dataset/helen/trainset/2960256451_1.jpg',
            '../../data/checkpoints-pfld-112/pfld-1983600' if not use_tflite else '../../data/pfld-112-quant.tflite',
            # '../../data/pfld-64.tflite',
            image_size=112)
    elif model == 'pfld-80':
        predict_single('/home/tamvm/Downloads/ibug_300W_large_face_landmark_dataset/helen/testset/3035796193_1_mirror.jpg', #'/home/tamvm/Downloads/ibug_300W_large_face_landmark_dataset/helen/trainset/2960256451_1.jpg',
            '../../data/checkpoints-pfld-80-025m/pfld-449100',
            # '../../data/pfld-64.tflite',
            zero_mean=False,
            depth_multiplier=0.25,
            image_size=80)
    elif model == 'pfld-custom-80':
        predict_single('/home/tamvm/Downloads/test_face_tamvm_2.jpg', #'/home/tamvm/Downloads/ibug_300W_large_face_landmark_dataset/helen/trainset/2960256451_1.jpg',
            '../../data/checkpoints-pfld-custom/pfld-183000',
            predict_fn=pfld_custom_predict_landmarks,
            # '../../data/pfld-64.tflite',
            depth_multiplier=1,
            zero_mean=True,
            image_size=80)
    elif model == 'pfld-custom-80-025m':
        predict_single('/home/tamvm/Downloads/test_face_tamvm_2.jpg', #'/home/tamvm/Downloads/ibug_300W_large_face_landmark_dataset/helen/trainset/2960256451_1.jpg',
            '../../data/checkpoints-pfld-custom-80-025m/pfld-314100',
            predict_fn=pfld_custom_predict_landmarks,
            # '../../data/pfld-64.tflite',
            depth_multiplier=0.25,
            zero_mean=True,
            image_size=80)
    elif model == 'pfld-custom-80-025m-aux7':
        predict_single('/home/tamvm/Downloads/test_face_tamvm_2.jpg', #'/home/tamvm/Downloads/ibug_300W_large_face_landmark_dataset/helen/trainset/2960256451_1.jpg',
            '../../data/checkpoints-pfld-custom-80-025m-aux7/pfld-376500',
            predict_fn=pfld_custom_predict_landmarks,
            # '../../data/pfld-64.tflite',
            depth_multiplier=0.25,
            zero_mean=True,
            image_size=80,
            aux_start_layer='layer_7')

    elif model == 'pfld-custom-80-025m-aux7-x3':
        predict_single( '/home/tamvm/Downloads/ibug_300W_large_face_landmark_dataset/helen/testset/3035796193_1_mirror.jpg', #'/home/tamvm/Downloads/ibug_300W_large_face_landmark_dataset/helen/trainset/2960256451_1.jpg',
            '../../data/checkpoints-pfld-custom-80-025m-aux7-x3/pfld-220000',
            predict_fn=pfld_custom_predict_landmarks,
            # '../../data/pfld-64.tflite',
            depth_multiplier=0.25,
            zero_mean=True,
            image_size=80,
            fc_x_n=3,
            box_detector='tf',
            aux_start_layer='layer_7')
    elif model == 'pfld-custom-80-025m-saux7-x3':
        predict_single( '/home/tamvm/Downloads/ibug_300W_large_face_landmark_dataset/helen/testset/3035796193_1_mirror.jpg', #'/home/tamvm/Downloads/ibug_300W_large_face_landmark_dataset/helen/trainset/2960256451_1.jpg',
            '../../data/checkpoints-pfld-custom-80-025m-saux7-x3/pfld-310500',
            predict_fn=pfld_custom_predict_landmarks,
            # '../../data/pfld-64.tflite',
            depth_multiplier=0.25,
            simple_aux=True,
            zero_mean=True,
            image_size=80,
            fc_x_n=3,
            box_detector='dlib',
            aux_start_layer='layer_7')
    elif model == 'pfld-custom-80-025m-aux7-x4-m3':
        predict_single('/home/tamvm/Downloads/test_face_tamvm_2.jpg',# '/home/tamvm/Downloads/ibug_300W_large_face_landmark_dataset/helen/testset/3035796193_1_mirror.jpg', #'/home/tamvm/Downloads/ibug_300W_large_face_landmark_dataset/helen/trainset/2960256451_1.jpg',
            '../../data/checkpoints-pfld-custom-80-025m-aux7-x4-m3/pfld-131500',
            predict_fn=pfld_custom_predict_landmarks,
            # '../../data/pfld-64.tflite',
            depth_multiplier=0.25,
            zero_mean=True,
            image_size=80,
            fc_x_n=4,
            mid_conv_n=3,
            box_detector='tf',
            aux_start_layer='layer_7')
    elif model == 'pfld-custom-80-025m-aux8':
        predict_single('/home/tamvm/Downloads/test_face_tamvm_2.jpg', #'/home/tamvm/Downloads/ibug_300W_large_face_landmark_dataset/helen/trainset/2960256451_1.jpg',
            '../../data/checkpoints-pfld-custom-80-025m-aux8/pfld-445500',
            predict_fn=pfld_custom_predict_landmarks,
            # '../../data/pfld-64.tflite',
            depth_multiplier=0.25,
            zero_mean=True,
            image_size=80,
            aux_start_layer='layer_8')
    else:
        use_tflite = True
        predict_single('/home/tamvm/Downloads/ibug_300W_large_face_landmark_dataset/helen/testset/3035796193_1_mirror.jpg', #'/home/tamvm/Downloads/ibug_300W_large_face_landmark_dataset/helen/trainset/2960256451_1.jpg',
            '../../data/landmark_80pose.tflite',
            normalize_lmks=True,
            # '../../data/pfld-64.tflite',
            image_size=80)


import click
from .core import NmapReportParser, NmapReport, CSVFileParser, JsonOutput, BateaModel, MatrixOutput
from defusedxml import ElementTree
from xml.etree.ElementTree import ParseError
from batea import build_report
import warnings
warnings.filterwarnings('ignore')

@click.command(context_settings=dict(help_option_names=['-h', '--help']))
@click.option("-c", "--read-csv", type=click.File('r'), multiple=True)
@click.option("-x", "--read-xml", type=click.File('r'), multiple=True)
@click.option("-n", "--n-output", type=int, default=5)
@click.option("-A", "--output-all", is_flag=True)
@click.option("-L", "--load-model", type=click.File('rb'), default=None)
@click.option("-D", "--dump-model", type=click.File('wb'), default=None)
@click.option("-f", "--input-format", type=str, default='xml')
@click.option('-v', '--verbose', count=True)
@click.option('-oM', "--output-matrix", type=click.File('w'), default=None)
@click.argument("nmap_reports", type=click.File('r'), nargs=-1)
def main(*, nmap_reports, input_format, dump_model, load_model,
         output_all, read_csv, read_xml, n_output, verbose, output_matrix):
    """Context-driven asset ranking based using anomaly detection"""

    report = build_report()
    csv_parser = CSVFileParser()
    xml_parser = NmapReportParser()
    if output_matrix:
        output_manager = MatrixOutput(output_matrix)
    else:
        output_manager = JsonOutput(verbose)

    try:
        if input_format == 'xml':
            for file in nmap_reports:
                report.hosts.extend([host for host in xml_parser.load_hosts(file)])
        if input_format == 'csv':
            for file in nmap_reports:
                report.hosts.extend([host for host in csv_parser.load_hosts(file)])
        if read_csv:
            for file in read_csv:
                report.hosts.extend([host for host in csv_parser.load_hosts(file)])
        if read_xml:
            for file in read_xml:
                report.hosts.extend([host for host in xml_parser.load_hosts(file)])
    except (ParseError, UnicodeDecodeError, ElementTree.ParseError, ValueError) as e:
        output_manager.log_parse_error(e)
        raise SystemExit

    if len(report.hosts) == 0:
        output_manager.log_empty_report()
        raise SystemExit

    report_features = report.get_feature_names()
    output_manager.add_report_info(report)

    matrix_rep = report.generate_matrix_representation()

    batea = BateaModel(report_features=report_features)

    if load_model is not None:
        batea.load_model(load_model)

    else:
        batea.build_model()
        batea.model.fit(matrix_rep)

    scores = -batea.model.score_samples(matrix_rep)
    output_manager.add_scores(scores)

    if output_all:
        n_output = len(scores)
    n_output = min(n_output, len(scores))

    top_n = scores.argsort()[-n_output:][::-1]

    for i, j in enumerate(top_n):
        output_manager.add_host_info(
            rank=str(i+1),
            score=scores[j],
            host=report.hosts[j],
            features={name: value for name, value in zip(report_features, matrix_rep[j, :])}
        )
    output_manager.flush()

    if dump_model:
        batea.dump_model(dump_model)

if __name__ == "__main__":
    main()


import numpy as np
import matplotlib.pyplot as plt
from shapely import geometry
from numpy.linalg import norm
from random import *
import pickle

def reach_set_calc(x_val, reach_range):
    """
    :type x_val: list
    :type reach_range: float
    :return: reach_set: Polygon
    Description: With given x and reach_range, generate a rectangular set centering at x with side length
                 2 * reach_range
    """
    p1 = geometry.Point(x_val[0] - reach_range, x_val[1] - reach_range)
    p2 = geometry.Point(x_val[0] + reach_range, x_val[1] - reach_range)
    p3 = geometry.Point(x_val[0] + reach_range, x_val[1] + reach_range)
    p4 = geometry.Point(x_val[0] - reach_range, x_val[1] + reach_range)
    vertex_list = [p1, p2, p3, p4]
    reach_set = geometry.Polygon(vertex_list)
    return reach_set

def all_Q_plt(Q, node_num, color_set, line_style_set, T, plt_scale):
    """
    :param Q: dict
    :param node_num: int
    :param color_set: list
    :param line_style_set: list
    :param T: int
    :return: None
    """
    # Plot all given convex sets
    for t_val in range(T + 1):
        for node in range(1, node_num + 1):
            hcoord_q, vcoord_q = Q[f"Q_t={t_val}^i={node}"].region.exterior.xy
            # plt.fill(hcoord_q, vcoord_q, alpha=0.1, facecolor=color_set[t_val], edgecolor=color_set[t_val],
            #          linewidth=2,
            #          linestyle=line_style_set[node - 1], label=fr"$Q_{t_val}^{{({node})}}$")
            plt.fill(hcoord_q, vcoord_q, alpha=1, facecolor='none', edgecolor=color_set[t_val],
                     linewidth=1.5,
                     linestyle=line_style_set[node - 1], label=r"$\mathcal{Q}_" + fr"{t_val}^{{({node})}}$")
    plt.legend(fontsize=14)
    plt.grid(True)
    plt.axis(plt_scale)
    return None

def set_plotter(set, plt_color, alpha_val):
    """
    :type set: Polygon
    :type plt_color: string
    :type alpha_val: float
    :return: None
    """
    hcoord, vcoord = set.exterior.xy
    # plt.fill(hcoord, vcoord, alpha=alpha_val, facecolor=plt_color, edgecolor=plt_color)
    plt.fill(hcoord, vcoord, alpha=alpha_val, facecolor='none', edgecolor=plt_color)

def game_plt(full_tree, oppo_action, Q, colors, UV_dict, t, prev_x_action, R, control):
    """
    :param full_tree: list
    :param oppo_action: State
    :param Q: dict
    :param colors: list
    :param UV_dict: dict
    :param t: int
    :param prev_x_action: State
    :param R: float
    :return: player_action: State
    """
    prev_x_state = prev_x_action.state
    ax = plt.gca()
    ax.set_aspect(1)

    # Plot selected Qt
    Qt = Q[f"Q_t={t}^i={oppo_action.state}"].region
    set_plotter(Qt, colors[t], alpha_val=0.05)

    # Plot the set discretized over
    if t == 0:
        set = Qt
    else:
        R_set = reach_set_calc(prev_x_action.state, R)
        set = Qt.intersection(R_set)
    set_plotter(set, colors[t], alpha_val=0.1)

    # Find disc xt in the set
    # disc_x_list = [action.state for action in full_tree if action.parent_state == oppo_action]

    # Plot disc xt in the set
    # for disc_x in disc_x_list:
    #     plt.scatter(disc_x[0], disc_x[1], color=colors[t], linewidths=0.1, marker='.')

    if control in ['1', '2']:   # Opt pl vs. Opt op or Opt pl vs. Sub-opt op
        # Find optimal player action xt
        player_action = UV_dict[f"V_t={t} ({prev_x_action.state}, {oppo_action.state})"].action
        player_state = player_action.state

    else:   # Control == '3', Sub-opt pl vs. Opt op
        # Randomly pick player action xt
        player_action = choice([action for action in full_tree if action.parent_state == oppo_action])
        player_state = player_action.state

    # Plot optimal xt in the set
    plt.scatter(player_state[0], player_state[1], color=colors[t], linewidths=1.5, marker='.')
    # plt.scatter(player_state[0], player_state[1], color='black', linewidths=0.1, marker='.')

    if t != 0:
        # Connect optimal xt state approximation to prev_x_state
        plt.plot([prev_x_state[0], player_state[0]], [prev_x_state[1], player_state[1]], color='black')
    return player_action

# Given a rectangular set, return discrete points inside the set
# def discrete_x_calc(poly, t_node, approx_para):
def discrete_x_calc(poly, approx_para, bound_rmv):
    """
    :type approx_para: int
    :type poly: Polygon
    :type bound_rmv: string
    :return discrete_x: list
    """
    [hcoord_val, vcoord_val] = poly.exterior.xy  # Find the horizontal and vertical coordinates of poly's vertices
    discrete_x = []
    for x_hcoord in np.linspace(min(hcoord_val), max(hcoord_val), approx_para):
        for x_vcoord in np.linspace(min(vcoord_val), max(vcoord_val), approx_para):
            discrete_x += [[x_hcoord, x_vcoord]]

    discrete_x_copy = discrete_x[:]  # Back up original discrete list
    if bound_rmv.lower() == 'y':
        # Find discrete x on the boundary
        bound_x = []
        for x_eval in discrete_x_copy:
            if x_eval[0] in hcoord_val or x_eval[1] in vcoord_val:
                bound_x.append(x_eval)

                # Remove discrete x on the boundary from original discrete list
                discrete_x.pop(discrete_x.index(x_eval))
        print(bound_x)
    return discrete_x

class State:
    # Set node loc i_t and parent_node i_t-1 as the attributes to newly defined OpNode object
    def __init__(self, state_value, parent_state, t_step, side):
        """
        :type state_value: int (Opponent), list (Player)
        :type parent_state: State / None (dummy_i)
        :type t_step: int, None(dummy i)
        :type side: str ('Player'/'Opponent')
        """
        self.state = state_value
        self.side = side
        self.parent_state = parent_state
        self.children_state_list = []
        self.t_step = t_step

    # Define methods that determine child nodes list with current node (Specific to graph)
    def add_child_state(self, child_state):
        """
        :type child_state: State
        """
        self.children_state_list.append(child_state)

class ConvexBody:
    def __init__(self, t_step, node, vertices):
        """
        :type t_step: int
        :type node: int
        :type vertices: list
        """
        self.t_step = t_step
        self.t_node = node
        self.region = geometry.Polygon(vertices)

class Value:  # Value function
    def __init__(self, player_state, oppo_state, t_step, side, value, action):
        """
        :type player_state: State (None for U0)
        :type oppo_state: State (None for U0)
        :type t_step: int
        :type side: string
        :type value: float
        :type action: State
        """
        self.side = side
        self.player_state = player_state
        self.oppo_state = oppo_state
        self.t_step = t_step
        self.value = value
        self.action = action

if __name__ == "__main__":
    #################################### Display ####################################
    """
    Plot trajectory result. Allow user to control opponent, while player always applies its optimal strategy by
    computer. Allow re-run functionality.
    IDEA: Separate game computation section (discretization, optimal value approximation) and game play section
    (display) as two different .py files
    """
    method = ''
    while method.lower() != 'terminate':
        method = input("Which method to use or terminate the program? [Old/New/Terminate]: ")
        if method.lower() == 'new':

            # Load tree info new files into the program
            tree_file = open(f'tree_info_new (epsilon = {0.15}, extra_disc_para = {5})', 'rb')
            tree_info = pickle.load(tree_file)
            tree_file.close()

            # Assign keys from tree_info to variables in this program
            Q = tree_info['Q']
            full_tree = tree_info['full_tree']
            UV_dict = tree_info['UV_dict']
            T = tree_info['T']
            num_nodes = tree_info['num_nodes']
            colors = tree_info['colors']
            line_style_list = tree_info['line_style_list']
            plt_scale = tree_info['plt_scale']
            extra_disc_para = tree_info['extra_disc_para']
            scale_para = tree_info['scale_para']
            dummy_i = tree_info['dummy_i']
            performance_bound = tree_info['performance_bound']
            R = tree_info['R']
            method = tree_info['method']

            # Plot all convex sets
            plt_scale_Q = [0, 0.8, 0, 0.8]
            all_Q_plt(Q, num_nodes, colors, line_style_list, T, plt_scale_Q)
            ax = plt.gca()
            ax.set_aspect(1)
            plt.show()
            plt_scale = [0.3, 0.4, 0.3, 0.4]

            msg = ''
            oppo_hist = dict()

            while msg.lower() != 'n':
                # Define figure and ax for result plot figure
                fig, ax = plt.subplots(figsize=(8, 8))
                for label in (ax.get_xticklabels() + ax.get_yticklabels()):
                    label.set_fontsize(22)

                tot_cost = 0
                all_Q_plt(Q, num_nodes, colors, line_style_list, T, plt_scale)
                ## Still need to add opt player vs. opt opponent
                control = input("Opt Player vs. Opt Opponent [1] / Opt Player vs. Sub-opt Opponent [2] / Sub-opt "
                                "Player vs. Opt Opponent [3]? ")
                if control not in ['1', '2', '3']:
                    print('Invalid game setting. Select again.')
                else:  # Valid game setting
                    if control == '2':  # Case of Player (PC) vs. Opponent (User)
                        # Initialize the game
                        t = 0
                        opt_player_action = dummy_i
                        opt_player_state = dummy_i.state

                        while t <= T:
                            prev_x_action = opt_player_action
                            prev_x_state = opt_player_state

                            oppo_node = int(input("Enter opponent action: "))
                            if t == 0:
                                if oppo_node not in range(num_nodes + 1):
                                    print("Invalid selection of node. Try again.")
                                else:  # oppo_node is valid with given graph
                                    oppo_hist[f"i{t}"] = oppo_node  # Store selected oppo_node to oppo_hist

                                    oppo_action = [action for action in full_tree if action.state == oppo_node and
                                                   action.parent_state == prev_x_action][0]

                                    # Plot the game process
                                    opt_player_action = game_plt(full_tree, oppo_action, Q, colors, UV_dict, t,
                                                                 prev_x_action, R, control)
                                    opt_player_state = opt_player_action.state
                                    # # Plot selected Q0
                                    # Q0 = Q[f"Q_t={t}^i={oppo_action.state}"].region
                                    # set_plotter(Q0, colors[t], alpha_val=0.25)
                                    # set_plotter(Q0, colors[t], alpha_val=0.5)
                                    #
                                    # # Find disc x0 in Q0
                                    # disc_x_list = [action.state for action in full_tree if action.parent_state ==
                                    #                oppo_action]
                                    #
                                    # # Plot disc x0 in Q0
                                    # for disc_x in disc_x_list:
                                    #     plt.scatter(disc_x[0], disc_x[1], color=colors[t], linewidths=0.5, marker='.')
                                    #
                                    # # Find optimal player action x0 (Can be made UDF)
                                    # opt_player_action = UV_dict[f"V_t={t} ({prev_x_action.state}, "
                                    #                             f"{oppo_action.state})"].action
                                    # opt_player_state = opt_player_action.state
                                    #
                                    # # Plot optimal x0 in Q0
                                    # plt.scatter(opt_player_state[0], opt_player_state[1], color='black', linewidths=0.1,
                                    #             marker='.')
                                    t += 1  # Update t value

                                    # Display
                                    print(f"Optimal Player State Approximation: {opt_player_state}")

                            else:  # t != 0
                                if oppo_node not in [action.state for action in prev_x_action.children_state_list]:
                                    print("Invalid selection of node. Try again.")
                                else:  # selected oppo_node is a reachable node
                                    oppo_hist[f"i{t}"] = oppo_node

                                    oppo_action = [action for action in full_tree if action.state == oppo_node and
                                                   action.parent_state == prev_x_action][0]

                                    # Plot the game process
                                    opt_player_action = game_plt(full_tree, oppo_action, Q, colors, UV_dict, t,
                                                                 prev_x_action, R, control)
                                    opt_player_state = opt_player_action.state
                                    # # Plot selected Qt
                                    # Qt = Q[f"Q_t={t}^i={oppo_action.state}"].region
                                    # set_plotter(Qt, colors[t], alpha_val=0.25)
                                    #
                                    # # Plot R(previous_x) intersect Qt
                                    # R_set = reach_set_calc(prev_x_state, R)
                                    # R_intersect_Q = Qt.intersection(R_set)
                                    # set_plotter(R_intersect_Q, colors[t], alpha_val=0.5)
                                    #
                                    # # Find disc xt in R(previous_x) intersect Qt
                                    # disc_x_list = [action.state for action in full_tree if action.parent_state ==
                                    #                oppo_action]
                                    #
                                    # # Plot disc xt in R(previous_x) intersect Qt
                                    # for disc_x in disc_x_list:
                                    #     plt.scatter(disc_x[0], disc_x[1], color=colors[t], linewidths=0.5, marker='.')
                                    #
                                    # # Find optimal player action xt in R(previous_x) intersect Qt
                                    # opt_player_action = UV_dict[f"V_t={t} ({prev_x_action.state}, {oppo_action.state}"
                                    #                             f")"].action
                                    # opt_player_state = opt_player_action.state
                                    #
                                    # # Plot optimal x_t in R(previous_x) intersect Qt
                                    # plt.scatter(opt_player_state[0], opt_player_state[1], color='black', linewidths=0.1,
                                    #             marker='.')
                                    #
                                    # # Connect optimal x_t state approximation to prev_x_state
                                    # plt.plot([prev_x_state[0], opt_player_state[0]],
                                    #          [prev_x_state[1], opt_player_state[1]], color='black')

                                    # Update cost
                                    tot_cost += norm(np.array(prev_x_state) - np.array(opt_player_state), 2)
                                    t += 1  # Update t value

                                    # Display
                                    print(f"Optimal Player State Approximation: {opt_player_state}")
                                    print(f"Cost: {tot_cost}")
                                    print(f"Running {method} method")

                        # plt.title(fr"Sub-optimal Opponent vs. Optimal Player " + '\n' +
                        #           fr"Opponent History: $i_0={oppo_hist['i0']}$, $i_1={oppo_hist['i1']}$, "
                        #           fr"$i_2={oppo_hist['i2']}$" + "\n" + fr"$\epsilon$={performance_bound}"
                        #           fr"(Without Boundary), Total Cost={round(tot_cost, 4)}")

                    elif control == '1':  # Case of Player (PC) vs. Opponent (PC)
                        for t in range(T+1):
                            if t == 0:
                                opt_oppo_action = UV_dict[f"U_t={t} ({dummy_i.state}, {None})"].action
                                prev_x_action = opt_oppo_action.parent_state

                                ### Need to check line 846 - 848 correctness!!!! Continue Here
                                # Plot game process
                                opt_player_action = game_plt(full_tree, opt_oppo_action, Q, colors, UV_dict, t,
                                                             prev_x_action, R, control)
                                opt_player_state = opt_player_action.state
                                prev_x_action = opt_player_action  # Reassign prev_x_action for next iteration use

                                # Update oppo_hist
                                oppo_hist[f"i{t}"] = opt_oppo_action.state

                            else:   # When t != 0
                                opt_oppo_action = UV_dict[f"U_t={t} ({prev_x_action.state}, {opt_oppo_action.state})"].\
                                    action

                                # Plot game process
                                opt_player_action = game_plt(full_tree, opt_oppo_action, Q, colors, UV_dict, t,
                                                             prev_x_action, R, control)
                                opt_player_state = opt_player_action.state

                                tot_cost += norm(np.array(prev_x_action.state) - np.array(opt_player_state), 2)
                                prev_x_action = opt_player_action

                                # Update oppo_hist
                                oppo_hist[f"i{t}"] = opt_oppo_action.state

                            # Display
                            print(f"\nt={t}")
                            print(f"Optimal i{t}: {opt_oppo_action.state}")
                            print(f"Optimal Player State Approximation: {opt_player_action.state}")
                            print(f"Total Cost: {tot_cost}")
                            print(f"Running {method} method")

                        # plt.title(fr"Optimal Opponent vs. Optimal Player " + '\n' +
                        #           fr"Opponent History: $i_0={oppo_hist['i0']}$, $i_1={oppo_hist['i1']}$, "
                        #           fr"$i_2={oppo_hist['i2']}$" + "\n" + fr"$\epsilon$={performance_bound}"
                        #           fr"(Without Boundary), Total Cost={round(tot_cost, 4)}")

                    elif control == '3':
                        for t in range(T+1):
                            if t == 0:
                                opt_oppo_action = UV_dict[f"U_t={t} ({dummy_i.state}, {None})"].action
                                prev_x_action = opt_oppo_action.parent_state

                                # Plot game process
                                ram_player_action = game_plt(full_tree, opt_oppo_action, Q, colors, UV_dict, t,
                                                             prev_x_action, R, control)
                                ram_player_state = ram_player_action.state
                                prev_x_action = ram_player_action

                                # Update oppo_hist
                                oppo_hist[f"i{t}"] = opt_oppo_action.state
                            else:
                                opt_oppo_action = UV_dict[f"U_t={t} ({prev_x_action.state}, {opt_oppo_action.state}"
                                                          f")"].action

                                ram_player_action = game_plt(full_tree, opt_oppo_action, Q, colors, UV_dict, t,
                                                             prev_x_action, R, control)
                                ram_player_state = ram_player_action.state

                                tot_cost += norm(np.array(prev_x_action.state) - np.array(ram_player_state), 2)
                                prev_x_action = ram_player_action

                                # Update oppo_hist
                                oppo_hist[f"i{t}"] = opt_oppo_action.state

                            # Display
                            print(f"\nt={t}")
                            print(f"Optimal i{t}: {opt_oppo_action.state}")
                            print(f"Sub-optimal Player State Approximation: {ram_player_state}")
                            print(f"Total Cost: {tot_cost}")
                            print(f"Running {method} method")

                        # plt.title(fr"Optimal Opponent vs. Sub-optimal Player " + '\n' +
                        #           fr"Opponent History: $i_0={oppo_hist['i0']}$, $i_1={oppo_hist['i1']}$, "
                        #           fr"$i_2={oppo_hist['i2']}$" + "\n" + fr"$\epsilon$={performance_bound}"
                        #           fr"(Without Boundary), Total Cost={round(tot_cost, 4)}")
                    plt.show()

                # Save Simulation Results
                sim_result = {
                    'tot_cost': tot_cost,
                    'performance_bound': performance_bound,
                    'extra_disc_para': extra_disc_para
                }
                sim_file = open(f'sim_result_new (epsilon = {performance_bound}, extra_disc_para = {extra_disc_para})',
                                'wb')
                pickle.dump(sim_result, sim_file)
                sim_file.close()

                msg = input(f"Rerun (Method: {method})? [Y/N] ")

            #################################### End Here ####################################

            #################################### Display ####################################
        elif method.lower() == 'old':
            tot_cost = 0
            # Load tree info old files into the program
            tree_file = open('tree_info_old', 'rb')
            tree_info = pickle.load(tree_file)
            tree_file.close()

            # Assign keys from tree_info to variables in this program
            Q = tree_info['Q']
            tree_no_lf_copy = tree_info['tree_no_lf_copy']
            UV_dict = tree_info['UV_dict']
            T = tree_info['T']
            num_nodes = tree_info['num_nodes']
            colors = tree_info['colors']
            line_style_list = tree_info['line_style_list']
            plt_scale = tree_info['plt_scale']
            disc_para = tree_info['disc_para']
            scale_para = tree_info['scale_para']
            dummy_i = tree_info['dummy_i']
            R = tree_info['R']
            method = tree_info['method']
            boundary_rmv = tree_info['boundary_rmv']

            msg = ''
            oppo_hist = dict()
            while msg.lower() != 'n':
                control = input("Player (PC) vs. Opponent (PC) [1] / Player (PC) vs. Opponent (User) [2]? ")
                if control not in ['1', '2']:
                    print('Invalid game setting. Select again.')
                else:  # control is in ['1', '2']
                    if control == '2':
                        # Let user be opponent, show player optimal action approximation for demo (Plot them)
                        t = 0
                        opt_player_action = None
                        opt_player_state = None
                        tot_cost = 0

                        while t <= T:
                            print(f"\nt={t}")

                            # I reassigned opt_player_action to avoid warning about potentially undefined
                            # opt_player_action in else statements
                            prev_x_action = opt_player_action
                            prev_x_state = opt_player_state  # Reassignment needed for later generation of R_intersect_Q

                            oppo_node = int(input("Enter opponent action: "))
                            if t == 0:
                                if oppo_node not in range(num_nodes + 1):
                                    print("Invalid selection of node. Try again.")
                                else:
                                    oppo_action = [action for action in tree_no_lf_copy if
                                                   action.state == oppo_node and action.t_step == t]
                                    oppo_action = oppo_action[0]

                                    oppo_hist[f"i{t}"] = oppo_action.state

                                    # Plot selected Q0
                                    Q0 = Q[f"Q_t={t}^i={oppo_node}"]
                                    set_plotter(Q0.region, colors[t], alpha_val=0.5)

                                    # Plot discrete x0 in Q0
                                    """
                                    disc_x0_list = [action.state for action in tree_no_lf_copy if
                                                    action.parent_state == oppo_action]
                                    """

                                    disc_x0_list = discrete_x_calc(Q[f'Q_t={t}^i={oppo_node}'].region, disc_para,
                                                                   bound_rmv=boundary_rmv)
                                    for disc_x0 in disc_x0_list:
                                        plt.scatter(disc_x0[0], disc_x0[1], color=colors[t], linewidths=0.5, marker='.')

                                    opt_player_action = UV_dict[
                                        f"V_t={t} ({oppo_action.parent_state.state}, {oppo_action.state})"].action
                                    opt_player_state = opt_player_action.state  # value of optimal x0 approximation

                                    print(f"Optimal Player State Approximation: {opt_player_state}")

                                    plt.scatter(opt_player_state[0], opt_player_state[1], color='black', linewidths=0.1,
                                                marker='.')
                                    t += 1
                            else:  # t != 0
                                if oppo_node not in [action.state for action in prev_x_action.children_state_list]:
                                    print("Invalid selection of node. Try again.")
                                else:
                                    oppo_action = \
                                        [state for state in tree_no_lf_copy if
                                         state.state == oppo_node and state.parent_state ==
                                         prev_x_action][0]

                                    oppo_hist[f"i{t}"] = oppo_action.state

                                    opt_player_action = UV_dict[
                                        f"V_t={t} ({prev_x_action.state}, {oppo_action.state})"].action
                                    opt_player_state = opt_player_action.state

                                    print(f"Optimal Player State Approximation: {opt_player_state}")

                                    # Plot Qt
                                    Qt = Q[f"Q_t={t}^i={oppo_action.state}"]
                                    set_plotter(Qt.region, colors[t], alpha_val=0.25)

                                    # Plot R(previous_x) intersect Q
                                    R_set = reach_set_calc(prev_x_state, R)
                                    R_intersect_Q = Q[f"Q_t={t}^i={oppo_action.state}"].region.intersection(R_set)
                                    set_plotter(R_intersect_Q, colors[t], alpha_val=0.5)

                                    # Plot discrete x in R_intersect_Q
                                    disc_x_list = discrete_x_calc(R_intersect_Q, approx_para=disc_para,
                                                                  bound_rmv=boundary_rmv)
                                    for disc_x in disc_x_list:
                                        plt.scatter(disc_x[0], disc_x[1], color=colors[t], linewidths=0.1, marker='.')
                                    # Plot optimal x_t state approximation
                                    plt.scatter(opt_player_state[0], opt_player_state[1], facecolor='black',
                                                linewidths=0.1, marker='.')
                                    # Connect optimal x_t state approximation to prev_x_state
                                    plt.plot([prev_x_state[0], opt_player_state[0]],
                                             [prev_x_state[1], opt_player_state[1]], color='black')

                                    tot_cost += norm(np.array(prev_x_state) - np.array(opt_player_state), 2)
                                    print(f"Total Cost: {tot_cost}")
                                    t += 1
                        if boundary_rmv.lower() == 'n':
                            plt.title(
                                fr"Opponent History: $i_0={oppo_hist['i0']}$, $i_1={oppo_hist['i1']}$, $i_2={oppo_hist['i2']}$ "
                                f"\n{disc_para}x{disc_para} Discretization, Total Cost={round(tot_cost, 4)}")
                        else:
                            plt.title(
                                fr"Opponent History: $i_0={oppo_hist['i0']}$, $i_1={oppo_hist['i1']}$, $i_2={oppo_hist['i2']}$ "
                                f"\n{disc_para}x{disc_para} Discretization (Without Boundary), Total Cost={round(tot_cost, 4)}")

                    elif control == '1':
                        opt_oppo_action = dummy_i
                        prev_x_action = opt_oppo_action.parent_state
                        tot_cost = 0
                        for t in range(T + 1):
                            if t == 0:
                                # Find optimal i_0
                                opt_oppo_action = UV_dict[f"U_t={t} ({dummy_i.state}, {None})"].action
                                # Plot Q0
                                Q0 = Q[f'Q_t={t}^i={opt_oppo_action.state}']
                                set_plotter(Q0.region, colors[t], alpha_val=0.25)
                                set_plotter(Q0.region, colors[t], alpha_val=0.5)

                                # Find discrete x0 in Q0
                                disc_x_list = [action.state for action in tree_no_lf_copy if action.parent_state ==
                                               opt_oppo_action]

                            else:  # when t is not 0
                                # Find optimal i_t
                                opt_oppo_action = UV_dict[
                                    f"U_t={t} ({prev_x_action.state}, {opt_oppo_action.state})"].action

                                # Plot selected Qt
                                Qt = Q[f"Q_t={t}^i={opt_oppo_action.state}"]
                                set_plotter(Qt.region, colors[t], alpha_val=0.25)

                                # Plot R(previous_x) intersect Q
                                R_set = reach_set_calc(prev_x_action.state, R)
                                R_intersect_Q = Qt.region.intersection(R_set)
                                set_plotter(R_intersect_Q, colors[t], alpha_val=0.5)

                                # Find discrete x in R_intersect_Q
                                disc_x_list = discrete_x_calc(R_intersect_Q, disc_para, bound_rmv=boundary_rmv)

                            # Output message
                            print(f"\nt={t}")
                            print(f"Optimal i{t}: {opt_oppo_action.state}")

                            # Plot discrete x in sets
                            for disc_x in disc_x_list:
                                plt.scatter(disc_x[0], disc_x[1], color=colors[t], linewidths=0.1, marker='.')

                            # Given x_t-1 and i_t, find approximation of optimal x_t
                            opt_player_action = \
                                UV_dict[
                                    f"V_t={t} ({opt_oppo_action.parent_state.state}, {opt_oppo_action.state})"].action

                            print(f"Optimal Player State Approximation: {opt_player_action.state}")

                            # Plot optimal x_t state approximation
                            plt.scatter(opt_player_action.state[0], opt_player_action.state[1], facecolor='black',
                                        linewidth=0.1,
                                        marker='.')

                            # Connect optimal x_t state approximation to prev_x_state
                            if t != 0:
                                plt.plot([prev_x_action.state[0], opt_player_action.state[0]],
                                         [prev_x_action.state[1], opt_player_action.state[1]], color='black')
                                # Update total cost
                                tot_cost += norm(np.array(prev_x_action.state) - np.array(opt_player_action.state), 2)
                                print(f"Total Cost: {tot_cost}")

                            prev_x_action = opt_player_action
                            # Store optimal opponent history
                            oppo_hist[f"i{t}"] = opt_oppo_action.state

                        # Plot display
                        if boundary_rmv.lower() == 'n':
                            plt.title(fr"Optimal Opponent History: $i_0={oppo_hist['i0']}$, $i_1={oppo_hist['i1']}$, "
                                      fr"$i_2={oppo_hist['i2']}$ "
                                      f"\n{disc_para}x{disc_para} Discretization, Total Cost={round(tot_cost, 4)}")
                        else:
                            plt.title(fr"Optimal Opponent History: $i_0={oppo_hist['i0']}$, $i_1={oppo_hist['i1']}$, "
                                      fr"$i_2={oppo_hist['i2']}$ "
                                      f"\n{disc_para}x{disc_para} Discretization (Without Boundary), Total Cost={round(tot_cost, 4)}")

                    # Plot all given convex sets
                    for t_val in range(T + 1):
                        for node in range(1, num_nodes + 1):
                            hcoord_q, vcoord_q = Q[f"Q_t={t_val}^i={node}"].region.exterior.xy
                            plt.fill(hcoord_q, vcoord_q, alpha=0.1, facecolor=colors[t_val], edgecolor=colors[t_val],
                                     linewidth=2,
                                     linestyle=line_style_list[node - 1], label=fr"$Q_{t_val}^{{({node})}}$")
                    plt.legend(fontsize=8)
                    plt.grid(True)
                    plt.axis(plt_scale)
                    if control == '1':
                        plt.savefig(f"Optimal Opponent History {oppo_hist['i0']}{oppo_hist['i1']}{oppo_hist['i2']}, "
                                    f"disc_para={disc_para}")
                    else:
                        plt.savefig(
                            f"Opponent History {oppo_hist['i0']}{oppo_hist['i1']}{oppo_hist['i2']}, disc_para={disc_para}")
                    plt.show()
                msg = input("Rerun? [Y/N] ")
            pass


#!/usr/bin/env python3
# coding=utf-8

'''MDMForm 系统配置主窗口'''

import os
import sys
from PyQt5 import QtCore, QtGui, QtWidgets
from PyQt5.QtGui import QPalette, QPixmap, QIcon
from PyQt5.QtWidgets import QMainWindow,QMessageBox,QTableWidgetItem,QFileDialog

from PyQt5 import QtSql
from PyQt5.QtSql import QSqlQuery

from openpyxl import load_workbook,Workbook

import warnings
warnings.filterwarnings('ignore')

BASE_DIR= os.path.dirname(os.path.dirname(os.path.abspath(__file__) ) )
sys.path.append( BASE_DIR  )
from ui.Ui_MDMForm import Ui_MDMForm

class MDMForm(QMainWindow,Ui_MDMForm):

    def __tablesql(self):
        """返回所有配置表的信息,用来对照Excel模板文件信息

        Returns:
            string: 查询表SQL语句
        """
        sql= 'SELECT object_id, \
                object_name, \
                object_name_cn, \
                object_desc, \
                template_file, \
                template_sheet, \
                start_row, \
                end_row \
                FROM xt_objects \
                where object_type=\'T\' \
                order by object_id ASC'
        return sql

    def __columnsql(self,id):
        """指定表的列及对应excel内汉字名称

        Args:
            id (int): 配置的表ID

        Returns:
            string: SQL查询列信息
        """
        sql= 'SELECT \
            object_name, \
            object_name_cn, \
            object_desc, \
            column_mapping \
            FROM xt_objects \
            where object_type=\'C\' \
            and parent_object_id= '
        sql += str(id)

        #sql += ' and rim(column_mapping) !=\'\' '
        sql += ' order by column_mapping asc'

        return sql

    def __columns(self,id):
        """指定表的列及对应excel内汉字名称

        Args:
            id (int): 配置的表ID

        Returns:
            dict: 回指定表的列及对应excel内汉字名称字典
        """
        cols = dict()
        if not self.db.isOpen():
            if not self.db.open():
                QMessageBox.critical(self, 'MDM', self.db.lastError().text())
                return
        query = QSqlQuery()
        if not query.exec(self.__columnsql(str(id))):
            QMessageBox.critical(self,'MDM', query.lastError().text())
        else:
            while query.next():
                cols[query.value('object_name')] = query.value('object_name_cn')
        return cols

    def __columnmap(self,id):
        """指定表的列及对应excel内对应列

        Args:
            id (int): 配置的表ID

        Returns:
            dict:  回指定表的列及对应excel内对应excel的列
        """
        cols = dict()
        if not self.db.isOpen():
            if not self.db.open():
                QMessageBox.critical(self, 'MDM', self.db.lastError().text())
                return
        query = QSqlQuery()
        if not query.exec(self.__columnsql(str(id))):
            QMessageBox.critical(self,'MDM', query.lastError().text())
        else:
            while query.next():
                cols[query.value('object_name')] = query.value('column_mapping')
        return cols

    def __insertsql(self,columnmap,tbname):
        """插入配置表语句

        Args:
            columnmap (dict): 配置表字段与Excel对应列
            tbname (string): 配置表名

        Returns:
            string: 带参数的sql insert语句
        """
        sqlk = ''
        sqlv = ''
        for k,v in columnmap.items():
            if len(str(v))>0:
                if len(sqlk) > 0:
                    sqlk +=','
                    sqlv +=','
                sqlk += str(k)
                sqlv += '?'
        sql = 'insert into '
        sql += tbname
        sql += ' ('
        sql += sqlk
        sql += ') values ('
        sql += sqlv
        sql +=')'
        return sql

    def __init__(self):
        super().__init__()
        self.setupUi(self)
        self.setupUiEx()
        self.addConnect()
        self.db = QtSql.QSqlDatabase.addDatabase('QSQLITE')
        self.db.setDatabaseName(os.path.join(BASE_DIR,'db\\mdm.db'))
        self.initData()
        if self.MDMListWidget.count()>0:
            self.MDMListWidget.setCurrentItem(self.MDMListWidget.item(0))
            self.mdmListClick()

    def closeEvent(self, QCloseEvent):
        if self.db.isOpen:
            self.db.close()

    def setupUiEx(self):
        palette = QPalette()
        icon = QIcon()
        appPath=os.path.join(BASE_DIR,u'res\\icon\\mdmconf.ico')
        icon.addPixmap(QPixmap(appPath))
        self.setWindowIcon(icon)

    def addConnect(self):
        self.MDMListWidget.clicked.connect(self.mdmListClick)
        self.btnTemplate.clicked.connect(self.templateClick)
        self.btnImport.clicked.connect(self.importClick)
        self.btnExport.clicked.connect(self.exportClick)
        self.btnUpdate.clicked.connect(self.updateClick)

    def initData(self):
        self.MDMListWidget.clear()
        if not self.db.isOpen():
            if not self.db.open():
                QMessageBox.critical(self, 'MDM', self.db.lastError().text())
                return
        query = QSqlQuery()
        if query.exec(self.__tablesql()):
            while query.next():
                qItem = QtWidgets.QListWidgetItem()
                cols = dict()
                cols['object_id'] = query.value('object_id')
                cols['object_name'] = query.value('object_name')
                cols['object_name_cn'] = query.value('object_name_cn')
                cols['object_desc'] = query.value('object_desc')
                cols['template_file'] = query.value('template_file')
                cols['template_sheet'] = query.value('template_sheet')
                cols['start_row'] = query.value('start_row')
                cols['end_row'] = query.value('end_row')
                qItem.setData(QtCore.Qt.ItemDataRole.UserRole,cols)
                qItem.setText(query.value('object_name_cn'))

                self.MDMListWidget.addItem(qItem)
        else:
             QMessageBox.critical(self,'MDM', query.lastError().text())

    def showData(self,id,name):
        """绑定配置数据到界面

        Args:
            id (int): 配置表ID
            name (string): 配置表名
        """
        cols = self.__columns(str(id))
        self.dataTableWidget.clear()
        self.dataTableWidget.setRowCount(0)
        self.dataTableWidget.setColumnCount(len(cols))
        self.dataTableWidget.setHorizontalHeaderLabels(cols.values())

        sql = ''
        for col in cols.keys():
            if(len(sql))>0:
                sql += ','
            sql += str(col)
        sql = 'select ' + sql
        sql +=' from '
        sql += str(name)
        if not self.db.isOpen():
            if not self.db.open():
                QMessageBox.critical(self, 'MDM', self.db.lastError().text())
                return
        query = QSqlQuery()
        if not query.exec(sql):
            QMessageBox.critical(self,'MDM', query.lastError().text())
        else:
            while query.next():
                rows=self.dataTableWidget.rowCount()
                self.dataTableWidget.insertRow(rows)
                for i in range(len(cols)):
                    qtitem=QTableWidgetItem(str(query.value(list(cols.keys())[i])))
                    self.dataTableWidget.setItem(rows,i,qtitem)

    def mdmListClick(self):
        qItem=self.MDMListWidget.currentItem()
        tconfs = dict(qItem.data(QtCore.Qt.ItemDataRole.UserRole))

        self.dataLabel.setText(str(tconfs['object_name_cn']) + ' : ' + str(tconfs['object_desc']))
        self.temFile.setText(str(tconfs['template_file']))
        self.temSheet.setText(str(tconfs['template_sheet']))
        self.temStart.setText(str(tconfs['start_row']))
        self.temEnd.setText(str(tconfs['end_row']))
        self.showData(str(tconfs['object_id']),str(tconfs['object_name']))

    def templateClick(self):
        if self.MDMListWidget.count()<=0:
            QMessageBox.information(self,'MDM', '请先选择要打开配置文件对应的基础数据配置表')
            return
        qItem=self.MDMListWidget.currentItem()
        tconfs = dict(qItem.data(QtCore.Qt.ItemDataRole.UserRole))

        if str(tconfs['template_file']) =='':
            QMessageBox.information(self,'MDM', '当前基础数据表尚未配置对应的配置文件')
            return
        appPath=os.path.join(BASE_DIR,str(tconfs['template_file']))
        #subprocess.run(appPath)
        os.system('start ' + appPath)
        #os.startfile(appPath)
        return

    def importClick(self):
        if self.MDMListWidget.count()<=0:
            QMessageBox.information(self,'MDM', '请先选择要重新导入数据的基础数据配置表')
            return
        fNames= QFileDialog.getOpenFileName(self,'导入基础数据', '/','Excel File (*.xlsx)')
        if not fNames[0]:
            return

        qItem=self.MDMListWidget.currentItem()
        tconfs = dict(qItem.data(QtCore.Qt.ItemDataRole.UserRole))
        sheetName = str(tconfs['template_sheet'])
        if QMessageBox.question(self, 'MDM', '确认更新模板配置表[' +sheetName + ']的数据?',QMessageBox.Yes|QMessageBox.No) == QMessageBox.No:
            return

        startRow = 2  #默认没有设置起始值，则默认从第二行开始（第一行为标题）
        if str(tconfs['start_row']).isdigit():
            startRow=int(str(tconfs['start_row']))
        endRow = 0   #没有设置结束行，默认后面数据行全部加载
        if str(tconfs['end_row']).isdigit():
            endRow=int(str(tconfs['end_row']))
        columnMap =dict()
        columnMap = self.__columnmap(str(tconfs['object_id']))

        try:
            wb= load_workbook(filename=fNames[0],read_only=True,data_only=True)
            if not (wb.sheetnames.index(sheetName) >= 0):
                QMessageBox.warning(self,'MDM', '选择的文件:' + fNames[0] + ',未包含配置指定的Sheet[' +sheetName + ']')
                wb.close()
                return

            ws=wb[sheetName]

            if endRow == 0:
                endRow = ws.max_row   # type: ignore
            sql = 'delete from ' + str(tconfs['object_name'])
            if not self.db.isOpen():
                if not self.db.open():
                    QMessageBox.critical(self, 'MDM', self.db.lastError().text())
                    return
            query = QSqlQuery()
            if not query.exec_(sql):
                QMessageBox.warning(self,'MDM', '清空数据表[' + str(tconfs['object_name_cn'] + ':' + query.lastQuery() + ']失败' + query.lastError().text()))
                wb.close()
                return
            sql = self.__insertsql(columnMap,str(tconfs['object_name']) )
            query.prepare(sql)
            for iRow in (range(startRow,endRow+1)):
                bAllEmptyflag = True
                for k,v in columnMap.items():
                    if len(str(v))<=0:
                        continue
                    if ws[str(v)+str(iRow)].value is None:
                        qvalue =''
                    else:
                        qvalue = str(ws[str(v)+str(iRow)].value)# type: ignore

                    if not(len(qvalue)==0 or qvalue.isspace()):
                        bAllEmptyflag = False
                    query.addBindValue(qvalue)

                if bAllEmptyflag:
                    continue
                elif not query.exec():
                        QMessageBox.warning(self,'MDM', '执行语句[' + query.lastQuery() + ']失败,' + query.lastError().text())
                        wb.close()
                        return
            wb.close()
            self.showData(str(tconfs['object_id']),str(tconfs['object_name']))
            QMessageBox.information(self,'MDM', '导入数据[' + str(tconfs['object_name_cn'])+ ']完成')
        except (NameError,ZeroDivisionError):
            QMessageBox.critical(self, '动力电缆计算', '变量名错误或除数为0')
        except OSError as reason:
            QMessageBox.critical(self, '动力电缆计算', str(reason))
        except TypeError as reason:
            QMessageBox.critical(self, '动力电缆计算', str(reason))
        except :
            QMessageBox.information(self,'动力电缆计算','导出数据文件失败')

    def exportClick(self):
        if self.MDMListWidget.count()<=0:
            QMessageBox.information(self,'MDM', '请先选择要重新导入数据的基础数据配置表')
            return
        fNames= QFileDialog.getSaveFileName(self,'下载基础数据', '/','Excel File (*.xlsx)')
        if not fNames[0]:
            return
        qItem=self.MDMListWidget.currentItem()
        tconfs = dict(qItem.data(QtCore.Qt.ItemDataRole.UserRole))
        sheetName = str(tconfs['template_sheet'])
        startRow = 2  #默认没有设置起始值，则默认从第二行开始（第一行为标题）
        if str(tconfs['start_row']).isdigit():
            startRow=int(str(tconfs['start_row']))

        columnMap = dict()
        column = dict()
        columnMap = self.__columnmap(str(tconfs['object_id']))
        column = self.__columns(str(tconfs['object_id']))

        try:
            wb = Workbook()
            ws = wb.active
            ws.title = sheetName
            for k,v in columnMap.items():
                if len(str(v))<=0:
                    continue
                ws[str(v)+str(startRow-1)] = column[k]   # type: ignore
            sql = ''
            for col in columnMap.keys():
                if(len(sql))>0:
                    sql += ','
                sql += str(col)
            sql = 'select ' + sql
            sql +=' from '
            sql += str(tconfs['object_name'])
            if not self.db.isOpen():
                if not self.db.open():
                    QMessageBox.critical(self, 'MDM', self.db.lastError().text())
                    return
            query = QSqlQuery()
            if not query.exec(sql):
                QMessageBox.critical(self,'MDM', query.lastError().text())
                return
            iRow = startRow
            while query.next():
                for k,v in columnMap.items():
                    if len(str(v))<=0:
                        continue
                    ws[str(v)+str(iRow)] = str(query.value(str(k))) # type: ignore
                iRow += 1
            wb.save(fNames[0])
            wb.close
            QMessageBox.information(self,'MDM','导出数据完成，文件名：' + fNames[0])
        except (NameError,ZeroDivisionError):
            QMessageBox.critical(self, '动力电缆计算', '变量名错误或除数为0')
        except OSError as reason:
            QMessageBox.critical(self, '动力电缆计算', str(reason))
        except TypeError as reason:
            QMessageBox.critical(self, '动力电缆计算', str(reason))
        except :
            QMessageBox.information(self,'动力电缆计算','导出数据文件失败')
        return

    def updateClick(self):
        if self.MDMListWidget.count()<=0:
            QMessageBox.information(self,'MDM', '请先选择要更新模板文件数据的基础数据配置表')
            return

        qItem=self.MDMListWidget.currentItem()
        tconfs = dict(qItem.data(QtCore.Qt.ItemDataRole.UserRole))
        if str(tconfs['template_file']) =='':
            QMessageBox.information(self,'MDM', '当前基础数据表尚未配置对应的配置文件')
            return
        tempfile=os.path.join(BASE_DIR,str(tconfs['template_file']))

        sheetName = str(tconfs['template_sheet'])

        if QMessageBox.question(self, 'MDM', '确认更新本地模板文件:' + tempfile + ',配置表[' + sheetName + ']的数据?',QMessageBox.Yes|QMessageBox.No) == QMessageBox.No:
            return

        startRow = 2  #默认没有设置起始值，则默认从第二行开始（第一行为标题）
        if str(tconfs['start_row']).isdigit():
            startRow=int(str(tconfs['start_row']))
        try:
            columnMap = dict()
            column = dict()
            columnMap = self.__columnmap(str(tconfs['object_id']))
            column = self.__columns(str(tconfs['object_id']))
            wb = load_workbook(tempfile,False)
            if not (wb.sheetnames.index(sheetName) >= 0):
                QMessageBox.warning(self,'MDM', '选择的文件:' + tempfile + ',未包含配置指定的Sheet[' + sheetName + ']')
                wb.close()
                return
            ws=wb[sheetName]
            #maxRow = ws.max_row # type: ignore
            # #暂未实现清除文档中老数据（考虑有附加列未导入数据库，如图片等）
            '''
            if startRow >1:
                for k,v in columnMap.items():
                    if len(str(v))<=0:
                        continue
                    QMessageBox.information(self,'MDM', str(column[k]))
                    ws[str(v)+str(startRow-1)] = str(column[k])   # type: ignore
                    #更新标题暂未实现
            '''
            sql = ''
            for col in columnMap.keys():
                if(len(sql))>0:
                    sql += ','
                sql += str(col)
            sql = 'select ' + sql
            sql +=' from '
            sql += str(tconfs['object_name'])
            if not self.db.isOpen():
                if not self.db.open():
                    QMessageBox.critical(self, 'MDM', self.db.lastError().text())
                    return
            query = QSqlQuery()
            if not query.exec(sql):
                QMessageBox.critical(self,'MDM', query.lastError().text())
                return
            iRow = startRow
            while query.next():
                for k,v in columnMap.items():
                    if len(str(v))<=0:
                        continue
                    ws[str(v)+str(iRow)] = str(query.value(str(k))) # type: ignore
                iRow += 1
            wb.save(tempfile)
            wb.close
            QMessageBox.information(self,'MDM','更新模板文件数据:' + tempfile + '完成')
        except (NameError,ZeroDivisionError):
            QMessageBox.critical(self, '动力电缆计算', '变量名错误或除数为0')
        except OSError as reason:
            QMessageBox.critical(self, '动力电缆计算', str(reason))
        except TypeError as reason:
            QMessageBox.critical(self, '动力电缆计算', str(reason))
        except :
            QMessageBox.information(self,'动力电缆计算','导出数据文件失败')
        return


import torch
import drjit as dr
import mitsuba as mi
import sys,os,json
import importlib
sys.path.append(".")
import cv2
import numpy as np
if torch.cuda.is_available():
    device = torch.device("cuda:0")
    torch.cuda.set_device(device)
else:
    device = torch.device("cpu")
from utils.logger import Logger
from utils.matcher import Matcher
from mitsuba.scalar_rgb import Transform4f as T
from tqdm.std import tqdm
mi.set_variant('cuda_ad_rgb')

log_level = 1

Pooler = torch.nn.AvgPool2d(kernel_size=2)
@dr.wrap_ad(source='drjit', target='torch')
def down_res_loss(st, img, img_ref):
    img = img[None,...].permute(0,3,1,2)
    img_ref = img_ref[None,...].permute(0,3,1,2)
    while st>0:
        img = Pooler(img)
        img_ref = Pooler(img_ref)
        st = st-1
    if log_level>0:
        Logger.save_img("down_res.png",img.permute(0,2,3,1)[0])
    return torch.mean((img-img_ref)**2)

if __name__=="__main__":

    method = sys.argv[1]
    config = sys.argv[2]
    Logger.init(exp_name=config+"/"+method, show=False, debug=False, path="results/",add_time=False)

    tasks = importlib.import_module(f'exp.{config}') # import specific task
    resolution = tasks.resolution #resolution
    spp = tasks.spp # spp
    scene = tasks.scene # scene
    thres = tasks.thres # for hybrid scheme
    max_depth = tasks.max_depth
    match_res = tasks.match_res

    # get target image
    if hasattr(tasks,"gt_img")==True:
        gt_img = torch.from_numpy(cv2.cvtColor(cv2.imread(tasks.gt_img),cv2.COLOR_BGR2RGB)).to(device)/255.0
        img_ref = mi.TensorXf(gt_img.reshape(-1,3))
    else:
        if hasattr(tasks,"gt_scene")==True:
            img_ref = mi.render(tasks.gt_scene, seed=0, spp=8192, sensor=0)
        else:
            img_ref = mi.render(scene, seed=0, spp=8192, sensor=0)
        img_ref = img_ref[...,:3]
        img_np = np.array(mi.util.convert_to_bitmap(img_ref))
        gt_img = torch.from_numpy(img_np).to(device)/255.0

    if log_level>0:
        Logger.save_img("gt_img.png",gt_img)

    gt_img_low= torch.from_numpy(cv2.resize(np.array(mi.util.convert_to_bitmap(img_ref)),(match_res,match_res))).to(device)/255.0

    # pixel matcher using optimal transport(Sinkhorn)
    matcher = Matcher(match_res, device)

    # get optimized parameter and transformation
    opt, apply_transformation, output, params = tasks.optim_settings()
    apply_transformation(params, opt)

    for key in opt.keys():
        dr.enable_grad(opt[key])
    params = mi.traverse(scene)

    # get init image
    img_init = mi.render(scene, params, seed=0, spp=512, sensor=0)
    init_img = torch.from_numpy(np.array( mi.util.convert_to_bitmap(img_init[...,:3]))).to(device)/255.0

    if log_level>0:
        Logger.save_img("init_img.png",init_img)

    # deal with hybrid scheme
    if method.endswith("hybrid"):
        method = method[:-7]
        integrator2 = mi.load_dict({
            'type': "prb_reparam",
            'max_depth': max_depth
        })
    else:
        thres = 10000

    # define integrator
    integrator1 = mi.load_dict({
        'type': method,
        'max_depth': max_depth
    })

    # camera settings are slightly different between EPSM and PRB.
    if method.startswith("manifold"):
        sensor_id = 1
    else:
        sensor_id = 0

    loop = tqdm(range(tasks.it))

    for it in loop:

        apply_transformation(params, opt)
        if it<thres:
            img = mi.render(scene, params, seed=it, spp=spp, integrator=integrator1, sensor=sensor_id)
        else:
            if it==thres:
                for key in opt.keys():
                    opt.reset(key)
            img = mi.render(scene, params, seed=it, spp=spp, integrator=integrator2, sensor=0)

        imgs = np.array(mi.util.convert_to_bitmap(img[...,:3]))

        if log_level>0:
            Logger.save_img(f"optim.png",imgs/255.0,flip=False)
            Logger.add_image(f"optim",imgs/255.0,flip=False)
        if log_level>1:
            Logger.save_img_2(f"optim{it}.png",imgs/255.0,flip=False)

        if img.shape[-1]==5:
            render_img = torch.from_numpy(cv2.resize(imgs,(match_res,match_res))).to(device)/255.0
            grad_ = matcher.match_Sinkhorn(render_img[...,:3].reshape(-1,3), gt_img_low[...,:3].reshape(-1,3))
            grad_ = grad_.reshape(match_res,match_res,5)
            grad_ = grad_.repeat(resolution//match_res,resolution//match_res,1)
            grad = mi.TensorXf(grad_)
            dr.backward(img*grad)
        else:
            # whether using multi-resolution loss
            # loss = down_res_loss(6-((7*it)//tasks.it),img,img_ref[...,:3])
            loss = dr.sum(dr.sqr(img - img_ref[...,:3])) / len(img)
            dr.backward(loss)

        try:
            # remove nan in grad
            dic = {}
            for key in opt.keys():
                x = dr.grad(opt[key])
                x[dr.isnan(x)] = 0
                dr.set_grad(opt[key],x)
                dic[key] = float(opt[key].torch().item())#.item()
            if log_level>1:
                Logger.save_param(f"param{it}.npy",dic)
        except:
            pass

        opt.step()
        loop.set_description(f"Iteration {it:02d}: error={output(opt)}")

    Logger.exit()
    img_final = mi.render(scene, params, seed=0, spp=8192, sensor=0)
    img_final = torch.from_numpy(np.array( mi.util.convert_to_bitmap(img_final[...,:3]))).to(device)/255.0
    if log_level>0:
        Logger.save_img(f"{sys.argv[1]}.png",img_final)

    print("finish optim")