import math

from six import iteritems

PARAM_K1 = 1.5
PARAM_B = 0.75
EPSILON = 0.25

class BM25(object):
    """Implementation of Best Matching 25 ranking function.

    Attributes
    ----------
    corpus_size : int
        Size of corpus (number of documents).
    avgdl : float
        Average length of document in `corpus`.
    doc_freqs : list of dicts of int
        Dictionary with terms frequencies for each document in `corpus`. Words used as keys and frequencies as values.
    idf : dict
        Dictionary with inversed documents frequencies for whole `corpus`. Words used as keys and frequencies as values.
    doc_len : list of int
        List of document lengths.
    """

    def __init__(self, corpus):
        """
        Parameters
        ----------
        corpus : list of list of str
            Given corpus.

        """
        self.corpus_size = 0
        self.corpus = corpus
        self.avgdl = 0
        self.doc_freqs = []
        self.idf = {}
        self.doc_len = []
        self.nd = {}
        self._initialize(corpus)

    def _initialize(self, corpus):
        """Calculates frequencies of terms in documents and in corpus. Also computes inverse document frequencies."""
        nd = {}  # word -> number of documents with word
        num_doc = 0
        for document in corpus:
            self.corpus_size += 1
            self.doc_len.append(len(document))
            num_doc += len(document)

            frequencies = {}
            for word in document:
                if word not in frequencies:
                    frequencies[word] = 0
                frequencies[word] += 1
            self.doc_freqs.append(frequencies)

            for word, freq in iteritems(frequencies):
                if word not in nd:
                    nd[word] = 0
                nd[word] += 1

        self.avgdl = float(num_doc) / self.corpus_size
        # collect idf sum to calculate an average idf for epsilon value
        idf_sum = 0
        # collect words with negative idf to set them a special epsilon value.
        # idf can be negative if word is contained in more than half of documents
        negative_idfs = []
        for word, freq in iteritems(nd):
            idf = math.log(self.corpus_size - freq + 0.5) - math.log(freq + 0.5)
            self.idf[word] = idf
            idf_sum += idf
            if idf < 0:
                negative_idfs.append(word)
        self.average_idf = float(idf_sum) / len(self.idf)

        eps = EPSILON * self.average_idf
        for word in negative_idfs:
            self.idf[word] = eps

    def get_score(self, document, index):
        """Computes BM25 score of given `document` in relation to item of corpus selected by `index`.

        Parameters
        ----------
        document : list of str
            Document to be scored.
        index : int
            Index of document in corpus selected to score with `document`.

        Returns
        -------
        float
            BM25 score.

        """
        score = 0
        doc_freqs = self.doc_freqs[index]
        for word in document:
            if word not in doc_freqs:
                continue
            score += (self.idf[word] * doc_freqs[word] * (PARAM_K1 + 1)
                      / (doc_freqs[word] + PARAM_K1 * (1 - PARAM_B + PARAM_B * self.doc_len[index] / self.avgdl)))
            # score += self.idf[word] * doc_freqs[word]

        return score

    def get_scores(self, document):
        """Computes and returns BM25 scores of given `document` in relation to
        every item in corpus.

        Parameters
        ----------
        document : list of str
            Document to be scored.

        Returns
        -------
        list of float
            BM25 scores.

        """
        scores = [self.get_score(document, index) for index in range(self.corpus_size)]
        return scores

    def get_words_score(self,document, index):

        words_score = {}
        doc_freqs = self.doc_freqs[index]
        for word in document:
            if word not in doc_freqs:
                continue
            score = (self.idf[word] * doc_freqs[word] * (PARAM_K1 + 1)
                      / (doc_freqs[word] + PARAM_K1 * (1 - PARAM_B + PARAM_B * self.doc_len[index] / self.avgdl)))
            # score = self.idf[word] * doc_freqs[word]
            if word not in words_score:
                words_score[word] = score
            else:
                words_score[word] = max(words_score[word],score)
        word_score_tuples = [(word, score) for word, score in words_score.items()]
        word_score_tuples = sorted(word_score_tuples,key=lambda x:x[1],reverse=True)
        return word_score_tuples

    def get_keywords(self, index, mention_tokens):
        document = self.corpus[index]
        intersection = list(set(document).intersection(set(mention_tokens)))
        words_scores = self.get_words_score(intersection,index)
        keywords = [ws[0] for ws in words_scores if ws[1] > 0 ]
        return keywords


import torch
import torchvision as tv
import torchvision.transforms as transforms
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import TensorDataset
from torch.autograd import Variable

from torch.nn import functional as F

import numpy as np
import matplotlib.pyplot as plt

from Input_initiation import input_initiation
from Input_selection import input_selection
from GP_build import conf_build, ratio_build, opc_predict

import warnings
warnings.filterwarnings('ignore')

import time
import os

import kernel_matrix

'''
Input selection and build gaussian model for c and r = c'/c
'''

seed = 1
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)

np.random.seed(1)

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

class Net(nn.Module):

    def __init__(self):
        super(Net, self).__init__()
        # Convolution (In LeNet-5, 32x32 images are given as input. Hence
        # padding of 2 is done below)
        self.conv1 = torch.nn.Conv2d(
            in_channels=1, out_channels=6, kernel_size=5, stride=1, padding=2, bias=True)
        # Max-pooling
        self.max_pool_1 = torch.nn.MaxPool2d(kernel_size=2)
        # Convolution
        self.conv2 = torch.nn.Conv2d(
            in_channels=6, out_channels=16, kernel_size=5, stride=1, padding=0, bias=True)
        # Max-pooling
        self.max_pool_2 = torch.nn.MaxPool2d(kernel_size=2)
        # Fully connected layer
        self.fc1 = torch.nn.Linear(16 * 2 * 2, 120)
        self.fc2 = torch.nn.Linear(120, 84)
        self.fc3 = torch.nn.Linear(84, 10)

    def forward(self, x):
        x = torch.nn.functional.relu(self.conv1(x))
        x = self.max_pool_1(x)
        x = torch.nn.functional.relu(self.conv2(x))
        x = self.max_pool_2(x)
        x = x.view(x.size()[0], -1)
        x = torch.nn.functional.relu(self.fc1(x))
        x = torch.nn.functional.relu(self.fc2(x))
        x = self.fc3(x)

        return x
    # last layer output

    def hidden(self, x):
        x = torch.nn.functional.relu(self.conv1(x))
        x = self.max_pool_1(x)
        x = torch.nn.functional.relu(self.conv2(x))
        x = self.max_pool_2(x)
        x = x.view(x.size()[0], -1)
        x = torch.nn.functional.relu(self.fc1(x))
        x = torch.nn.functional.relu(self.fc2(x))

        return x

def orig_profit(model, x_test, y_test):
    # get confidence
    softmaxes = F.softmax(model.fc3(x_test), dim=1)
    confidences, predictions = torch.max(softmaxes, 1)
    predictions = predictions.detach().numpy()
    confidences = confidences.detach().numpy()
    y_test = y_test.detach().numpy()

    index = np.where(predictions != y_test)
    corr = np.sum(confidences[index] > 0.9)
    print('high mis is ', corr)
    index = np.where(predictions == y_test)
    incorr = np.sum(confidences[index] > 0.9)
    print('high correct is ', incorr)

    deter = predictions == y_test
    score = np.square(deter - confidences)
    bs = np.mean(score)
    print('borel score is {}'.format(bs))
    acc = np.mean(deter)
    unc = acc * (1 - acc)
    print('uncertainy is {}'.format(unc))
    return corr, incorr, bs, unc

def calibrated_profit(model, clf, x_test, y_test, center_):
    softmaxes = F.softmax(model.fc3(x_test), dim=1)
    _, predictions = torch.max(softmaxes, 1)
    predictions = predictions.detach().numpy()
    y_test = y_test.detach().numpy()

    confidences, _ = opc_predict(model, clf, x_test, center_)
    indicate = predictions == y_test

    index = np.where(predictions != y_test)
    incorr = np.sum(confidences[index] > 0.9)
    # print('high mis is ', corr)
    index = np.where(predictions == y_test)
    corr = np.sum(confidences[index] > 0.9)
    # print('high correct is ', incorr)

    deter = predictions == y_test
    score = np.square(deter - confidences)
    bs = np.mean(score)
    # print('borel score is {}'.format(score))

    return corr, incorr, bs

def load_results(stat, init_size, final_size):
    orig_profit(model, x_test, y_test)

    corr, incorr, bs = 0, 0, 0
    for k in range(stat):
        ind = torch.load('exp_results/index-record/stat{0}.pt'.format(k))
        center_ind = ind[0:init_size]
        center_ = x_op[center_ind].detach().numpy()
        index = ind[0:final_size + init_size]
        x_select = x_op[index]
        y_select = y_op[index]
        select_index = index
        gp = ratio_build(model, x_op, x_select,
                         y_select, select_index, center_)

        c, i, b = calibrated_profit(model, gp, x_test, y_test, center_)
        corr += c
        incorr += i
        bs += b
    print('size is {0},high conf corr is {1}, high conf incorr is {2}, average bs is {3}'.format(
        final_size, corr / stat, incorr / stat, bs / stat))

if __name__ == '__main__':
    # load original model
    model_dir = 'model/LeNet-5.pth'
    model = Net()
    model.load_state_dict(torch.load(model_dir))

    # load operational test data
    x_test, y_test = torch.load('./data/test.pt')
    x_test = torch.FloatTensor(x_test)
    y_test = torch.LongTensor(y_test)
    x_test = torch.unsqueeze(x_test, dim=1)
    x_test, y_test = x_test.to(device), y_test.to(device)
    x_test = model.hidden(x_test)

    # load operational valid data
    x_valid, y_valid = torch.load('./data/valid.pt')
    x_valid = torch.FloatTensor(x_valid)
    y_valid = torch.LongTensor(y_valid)
    x_valid = torch.unsqueeze(x_valid, dim=1)
    x_op, y_op = x_valid.to(device), y_valid.to(device)
    x_op = model.hidden(x_op)

    stat = 8
    init_size = 10
    final_size = 600
    load_results(stat, init_size, final_size)


#!/bin/env python3
import itertools
import re
import sys
from os import path, walk

def main(argv):
    testnet_dir = path.dirname(path.abspath(argv[0]))
    env_dir = path.join(testnet_dir, "env")

    res = list(itertools.chain(*map(triples_from_file, host_files(env_dir))))

    testnets = sorted(set(m[0] for m in res))
    hosts = sorted(set(m[3] for m in res))

    table = {h: {tn: [] for tn in testnets} for h in hosts}

    for tn, typ, idx, host in res:
        table[host][tn].append("{}.{}".format(typ, idx))

    print("host;{}".format(";".join(sorted(testnets))))
    for k, v in sorted(table.items()):
        line = ";".join([", ".join(v[tn]) for tn in testnets])
        print("{};{}".format(k, line))

#    print(sorted(testnets))
#    print(len(list(testnets)))
#    print(sorted(hosts))
#    print(len(list(hosts)))
#    print(table)

def host_files(p):
    for (root, dirs, files) in walk(p):
        for f in filter(lambda f: f.endswith(".ini"), files):
            yield path.join(root, f)

NODE_DEF_RE = re.compile(r'^\s*([-a-zA-Z.0-9]+)\s+.*ic_host="([a-zA-Z-0-9.]+)"')
NODE_DEF_RE = re.compile(r'^\s*([-a-zA-Z0-9]+)\.([a-zA-Z0-9]+)\.([a-zA-Z0-9]+)\s+.*ic_host="([a-zA-Z-0-9.]+)"')

def triples_from_file(h):
    with open(h, "r") as f:
        matches = (NODE_DEF_RE.match(line) for line in f)
        matches = (m for m in matches if m is not None)
        return [(m[1], m[2], m[3], m[4]) for m in matches]

if __name__ == "__main__":
    main(sys.argv)


import time
import cv2
import math
import numpy as np
from Person import *

class Camera(object):
    def __init__(self, input):

        #clock
        self.trustworth_time = 10 #moing queue1 to querue2 every 10s
        self.sys_clear_time = self.trustworth_time/2
        self.no_detection_time = time.time()

        self.input = input
        self.video = cv2.VideoCapture(self.input)
        # self.video.set(3,1280)
        # self.video.set(4,720)
        self.w = int(self.video.get(3))
        self.h = int(self.video.get(4))
        self.rangeLeft = int(1*self.w/8)
        self.rangeRight = int(7*self.w/8)
        self.midLine = int(4*self.w/8)

        self.start_time = time.time()
        self.entered = 0

        #queue1
        self.persons= []
        self.pid = 1
        self.check_key = True
        self.last_checklist = {}

        #queue2
        self.stable_persons= []
        self.stable_pid = 1
        self.stable_check_key = True
        self.stable_last_checklist = {}
        self.display_pid = []

        #queue3
        self.valid_persons = []

    def __del__(self):
        self.video.release()

    def frameDetections(self):
        return self.video.read()

    def fack_preson_check(self):
        #check every ppl location in every "sys_clear_time"
        '''

        '''
        if int(time.time() - self.start_time) % self.sys_clear_time == 0 and self.check_key == True:
            cur_checklist = {}
            cur_indexlist = {}
            for index, p in enumerate(self.persons):

                fix_dist = math.sqrt((p.getX())**2 + (p.getY())**2)
                cur_checklist["pid" + str(p.getId())] = fix_dist
                cur_indexlist["pid" + str(p.getId())] = index

            self.check_key = False
            for key in self.last_checklist.keys():
                try:
                    if cur_checklist[key] == self.last_checklist[key]:
                        print("[POP] pop q1 : ", key)
                        self.persons.pop(cur_indexlist[key])
                        cur_checklist.pop(key)

                        ## DEBUG:
                        #print("[INFO] stay in q1 pid = " , [k for k in cur_checklist.keys()])

                except Exception as e:
                    pass

            self.last_checklist = cur_checklist

        elif int(time.time() - self.start_time) % self.sys_clear_time != 0:
            self.check_key = True

    def stabel_fack_preson_check(self):
        #check every ppl location in every sys_clear_time
        if int(time.time() - self.start_time) % self.sys_clear_time == 0 and self.stable_check_key == True:
            cur_checklist = {}
            cur_indexlist = {}
            for index, p in enumerate(self.stable_persons):
                fix_dist = math.sqrt((p.getX())**2 + (p.getY())**2)
                cur_checklist["pid" + str(p.getId())] = fix_dist
                cur_indexlist["pid" + str(p.getId())] = index

            self.stable_check_key = False

            for key in self.stable_last_checklist.keys():
                try:
                    if cur_checklist[key] == self.stable_last_checklist[key]:

                        print("[MOVE] Q2 pid{} -> Q3".format(p.getId()))
                        p = self.stable_persons[cur_indexlist[key]]
                        p.updateLeavetime(time.time()-self.sys_clear_time) #update person leave time
                        self.stable_persons.pop(cur_indexlist[key])
                        cur_checklist.pop(key)
                        self.face_pool.pop(key) #clear face pool
                        self.entered += 1
                        self.valid_persons.append(p)

                except Exception as e:
                    pass

            self.stable_last_checklist = cur_checklist

        elif int(time.time() - self.start_time) % self.sys_clear_time != 0:
            self.stable_check_key = True

    def people_tracking(self, rects):

        '''
        if no detection longer than 3min flush system
        '''
        # if len(rects) == 0 and (time.time() - self.no_detection_time) > 180:
        #     self.persons = []
        #     self.stable_persons = []
        #     self.no_detection_time = time.time()
        #     print("[WARNING] no detection")

        self.display_pid = [] #fresh display list must here !!!!!
        for person in rects:
            xCenter, yCenter, w, h = person['rect']
            gender = person['gender']
            age = person['age']
            proj = person['project']
            new = True
            inActiveZone= xCenter in range(self.rangeLeft,self.rangeRight)

            """
            queue1
            """
            self.fack_preson_check()

            for index, p in enumerate(self.persons):
                #if person stay in frame over 5s
                if time.time()- p.getEnter_t() > self.trustworth_time:
                    print("[MOVE] q1 pid{} -> q2 pid{} ".format(p.getId(), str(self.stable_pid)))
                    p.updatePid(self.stable_pid) #replace the pid with new pid
                    self.stable_persons.append(p)
                    self.persons.pop(index) #pop person from the persons list
                    self.stable_pid += 1

                dist = math.sqrt((xCenter - p.getX())**2 + (yCenter - p.getY())**2)
                p.updateAttris(age, gender, proj)

                if dist <= w and dist <=h:
                    if inActiveZone:
                        new = False
                        self.display_pid.append([p.getId(), (10, 10, 200)]) #id, display color
                        p.updateCoords(xCenter,yCenter)
                        break
                    else:
                        try:
                            self.persons.pop(index)
                            print("[POP]  pid{} removed from q1".format(str(p.getId())))
                            self.last_checklist.pop("pid" + str(p.getId())) #pop last frame dict id --- > 1st then pop list
                        except Exception as e:
                            pass

            """
            queue2
            """
            self.stabel_fack_preson_check()
            for index, p in enumerate(self.stable_persons):
                dist = math.sqrt((xCenter - p.getX())**2 + (yCenter - p.getY())**2)
                p.updateAttris(age, gender, proj)
                if dist <= w and dist <=h:
                    if inActiveZone:
                        new = False
                        self.display_pid.append([p.getId(), (10, 200, 10)]) #id, display color
                        p.updateCoords(xCenter,yCenter)
                        break
                    else:
                        self.entered += 1 #count ppl in the area
                        try:
                            self.stable_last_checklist.pop("pid" + str(p.getId())) #pop last frame dict id --- > 1st then pop list
                        except Exception as e:
                            pass

                        print("[MOVE] q2 pid{} -> q3".format(p.getId()))
                        self.stable_persons.pop(index)
                        p.updateLeavetime(time.time()) #update person leave time
                        self.valid_persons.append(p)
            #make sure the total persons number wont excess the bounding box
            if new == True and inActiveZone and len(self.persons) + len(self.stable_persons) + 1 <= len(rects) :
                print("[CREAT] new pid" + str(self.pid))
                enter_t = time.time()
                p = Person(self.pid, xCenter, yCenter, enter_t)
                self.persons.append(p)
                self.pid += 1


import logging
import os

import pkg_resources

from typing import Dict
from pathlib import Path

from flask.cli import load_dotenv
from sentry_sdk.integrations.flask import FlaskIntegration
from str2bool import str2bool

class Config:
    """
    Flask configuration base class

    Includes a mixture of static and dynamic configuration options. Dynamic objects are typically set from environment
    variables (set directly or through environment files).

    See the project README for configuration option details.
    """

    ENV = os.environ.get("FLASK_ENV")
    DEBUG = False
    TESTING = False

    NAME = "bas-web-map-inventory"

    LOGGING_LEVEL = logging.WARNING

    def __init__(self):
        load_dotenv()

        self.APP_ENABLE_FILE_LOGGING = str2bool(os.environ.get("APP_ENABLE_FILE_LOGGING")) or False
        self.APP_ENABLE_SENTRY = str2bool(os.environ.get("APP_ENABLE_SENTRY")) or True

        self.LOG_FORMAT = "[%(asctime)s] %(levelname)s [%(name)s.%(funcName)s:%(lineno)d] %(message)s"
        self.LOG_FILE_PATH = Path(os.environ.get("APP_LOG_FILE_PATH") or "/var/log/app/app.log")

        self.SENTRY_DSN = os.environ.get("SENTRY_DSN") or None

        self.AIRTABLE_API_KEY = os.environ.get("AIRTABLE_API_KEY")
        self.AIRTABLE_BASE_ID = os.environ.get("AIRTABLE_BASE_ID")

    # noinspection PyPep8Naming
    @property
    def VERSION(self) -> str:
        return "Unknown"

    # noinspection PyPep8Naming
    @property
    def SENTRY_CONFIG(self) -> Dict:
        return {
            "dsn": self.SENTRY_DSN,
            "integrations": [FlaskIntegration()],
            "environment": self.ENV,
            "release": f"{self.NAME}@{self.VERSION}",
        }

class ProductionConfig(Config):  # pragma: no cover
    """
    Flask configuration for Production environments

    Note: This method is excluded from test coverage as its meaning would be undermined.
    """

    def __init__(self):
        super().__init__()
        self.APP_ENABLE_FILE_LOGGING = str2bool(os.environ.get("APP_ENABLE_FILE_LOGGING")) or True

    # noinspection PyPep8Naming
    @property
    def VERSION(self) -> str:
        return pkg_resources.require("bas-web-map-inventory")[0].version

class DevelopmentConfig(Config):  # pragma: no cover
    """
    Flask configuration for (local) Development environments

    Note: This method is excluded from test coverage as its meaning would be undermined.
    """

    DEBUG = True

    @property
    def SENTRY_CONFIG(self) -> Dict:
        _config = super().SENTRY_CONFIG
        _config["server_name"] = "Local container"

        return _config

    LOGGING_LEVEL = logging.INFO

    def __init__(self):
        super().__init__()
        self.APP_ENABLE_SENTRY = str2bool(os.environ.get("APP_ENABLE_SENTRY")) or False

    # noinspection PyPep8Naming
    @property
    def VERSION(self) -> str:
        return "N/A"

class TestingConfig(Config):
    """
    Flask configuration for Testing environments
    """

    DEBUG = True
    TESTING = True

    LOGGING_LEVEL = logging.DEBUG

    def __init__(self):
        super().__init__()
        self.APP_ENABLE_FILE_LOGGING = False
        self.APP_ENABLE_SENTRY = False

    # noinspection PyPep8Naming
    @property
    def VERSION(self) -> str:
        return "N/A"


import os
import prody as prd
import math
from collections import defaultdict
import itertools
import pandas as pd
from multiprocessing import Pool

def distance(distances_1, distances_2):
    x1, y1, z1 = distances_1
    x2, y2, z2 = distances_2
    distance = math.sqrt(math.pow(x2 - x1, 2) +
               math.pow(y2 - y1, 2) +
               math.pow(z2 - z1, 2) * 1.0)
    return distance

def ss_bonds(pdb):

    try:
        cysteines = defaultdict(list)
        free_cys = 0
        cys_bridges = 0

        prot = prd.parsePDB(pdb, QUIET=True)
        for residue in prot.iterResidues():
            name_id = str(residue).split()
            name_id. insert(0, str(residue.getChain()).split()[1])
            if name_id[1] == 'CYS':
                cysteines['_'.join(name_id)].extend(list(residue.getAtom('SG').getCoords()))
                #print('_'.join(name_id))

        ss_bonds = []

        for pair in itertools.combinations(cysteines.keys(), r=2):
            if pair[0] != pair[1]:
                d = distance(cysteines[pair[0]], cysteines[pair[1]])
                if d <= 4.0:
                    ss_bonds.append((pair) + (d, ))
                    cys_bridges += 1
                else:
                    free_cys +=1

    except Exception:
        free_cys, cys_bridges = 0, 0

    return free_cys, cys_bridges

cpus = 70

path = '/storage/evagsm/nobackup/sensitivity_h/'
files = [path + f for f in os.listdir(path) if f.endswith('.pdb')]

pool = Pool(processes = cpus)
results = pool.map(ss_bonds, files)
pool.close()
pool.join()

print(len(results))

df = pd.DataFrame(results)
df.columns = ['free_cys', 'cys_bridges']
df.insert(0, 'ID', [s.split('/')[-1][:-4] for s in files])
df.to_csv('./sens_CYS.csv', index=None)


# pylint: skip-file
import json

import joblib
import numpy as np
from azureml.core.model import Model

from package.train_pipeline.constants import (
    TrainPipelineConstants
)

def init():

    global model

    model_path = Model.get_model_path(TrainPipelineConstants.MODEL_NAME.value)
    model = joblib.load(model_path)

def run(data):

    try:
        data = json.loads(data)
        data = data['data']
        result = model.predict(np.array(data))
        return result.tolist()
    except Exception as e:
        error = str(e)
        return error


#!/usr/bin/python3

from selenium import webdriver
from selenium.webdriver import ActionChains
from selenium.webdriver.common.keys import Keys
from selenium.webdriver.common.by import By
import time

opts = webdriver.FirefoxOptions()
opts.add_argument("--headless")
driver = webdriver.Firefox(firefox_options=opts)
driver.get("http://localhost/Identity/Login")

time.sleep(1)
u_input = driver.find_element_by_name("Data.Email")
p_input = driver.find_element_by_name("Data.Password")
u_input.send_keys("admin@bugznet.com")
p_input.send_keys("!bugzn3t@admin")
time.sleep(1)
p_input.send_keys(Keys.RETURN)
time.sleep(1)

driver.get("http://localhost/Bugs")
print("Admin user clicking around the ui")

links = driver.find_elements(By.CSS_SELECTOR, 'a[href*="/Bugs/"]')
link_count=len(links)
for i in range(link_count):
    links = driver.find_elements(By.CSS_SELECTOR, 'a[href*="/Bugs/"]')
    driver.execute_script("arguments[0].scrollIntoView(true);", links[i])
    time.sleep(1)
    links[i].click()
    time.sleep(1)
    form = driver.find_element_by_tag_name('form')
    action = ActionChains(driver)
    action.move_to_element_with_offset(form,50,50).click().perform()
    time.sleep(1)
    driver.back()
    time.sleep(1)

time.sleep(2)
driver.close()


import d2lzh as d2l
from mxnet import autograd, gluon, init, nd
from mxnet.gluon import nn, data as gdata, loss as gloss
import numpy as np
import time

def get_data_ch7():
    data = np.genfromtxt('data/airfoil_self_noise.dat', delimiter = '\t')
    data = (data - data.mean(axis=0)) / data.std(axis=0)
    return nd.array(data[:1500, : -1]), nd.array(data[:1500, -1])

features, labels = get_data_ch7()
print(features.shape)

def sgd(params, states, hyperparams):
    for p in params:
        p[:] -= hyperparams['lr'] * p.grad

def train_ch7(trainer_fn, states, hyperparams, features, labels, batch_size=10, num_epochs=2):
    # initialize model
    net, loss = d2l.linreg, d2l.squared_loss
    w = nd.random.normal(scale=0.01, shape=(features.shape[1], 1))
    b = nd.zeros(1)
    w.attach_grad()
    b.attach_grad()

    def eval_loss():
        return loss(net(features, w, b), labels).mean().asscalar()

    ls = [eval_loss()]
    data_iter = gdata.DataLoader(gdata.ArrayDataset(features, labels), batch_size, shuffle=True)
    for _ in range(num_epochs):
        start = time.time()
        for batch_i, (X, y) in enumerate(data_iter):
            with autograd.record():
                l = loss(net(X, w, b), y).mean() # use average loss
            l.backward()
            trainer_fn([w, b], states, hyperparams) # iter model params
            if(batch_i + 1) * batch_size % 100 == 0:
                ls.append(eval_loss()) # record the moment train loss per 100 samples

        # print ouput and draw
        print('loss: %f, %f sec per epoch' % (ls[-1], time.time() - start))
        d2l.set_figsize()
        d2l.plt.plot(np.linspace(0, num_epochs, len(ls)), ls)
        d2l.plt.xlabel('epoch')
        d2l.plt.ylabel('loss')

def train_sgd(lr, batch_size, num_epochs=2):
    train_ch7(sgd, None, {'lr':lr}, features, labels, batch_size, num_epochs)

# train_sgd(1, 1500, 6)
train_sgd(0.005, 1)


import json
# import re
import requests
from requests.structures import CaseInsensitiveDict
from date.date import ipaddress, headers_login

# ipaddress = "192.168.1.100"

def auth_token(ipaddress, headers_login):
    url = f"http://{ipaddress}/login"
    headers = CaseInsensitiveDict()
    headers["Content-Type"] = "application/json"
    data = headers_login
    auth = requests.post(url, headers=headers, data=data)
    token_auth = json.loads(auth.text)["sessionId"]
    return token_auth

def get_users(ipaddress, sessionId):
    url = f"http://{ipaddress}/api/settings/users/get_users"
    headers["sessionId"] = sessionId
    auth = requests.get(url, headers=headers)
    if auth.status_code == 200 :
        users = json.loads(auth.text)["get_users"]
        return users
    else:
        return None

def get_network_ip(ipaddress, sessionId):
    url = f"http://{ipaddress}/api/network/tcpip/ip"
    headers["sessionId"] = sessionId
    auth = requests.get(url, headers=headers)
    print(auth.status_code)
    print(auth.text)
    if auth.status_code == 200 :
        network_ip = json.loads(auth.text)
        return network_ip
    else:
        return None


from typing import Any, Tuple

import torch
import torch.distributions as D
import torch.nn as nn
import torch.nn.functional as F

from ..tools import *
from .actorcritic import *
from .networks.common import *
from .math_functions import *
from .networks.encoders import *
from .networks.decoders import *
from .networks.rssm_component import *
from .networks.rssm import *
# from .rssm_simplified import RSSMCore, RSSMCell
from .probes import *
from .tools_v3 import *
from .networks import *
from .world_models import *

class Dreamer_agent(nn.Module):

    def __init__(self, conf,obs_space,act_space,step,device=None):
        super().__init__()
        assert conf.action_dim > 0, "Need to set action_dim to match environment"
        features_dim = conf.deter_dim + conf.stoch_dim * (conf.stoch_discrete or 1)
        self.iwae_samples = conf.iwae_samples
        self.imag_horizon = conf.imag_horizon
        if device:
            self._device=device
        else:
            self._device=conf.device

        # World model
        self.wm_type=conf.wm_type
        # if self.wm_type=='v2':
        self.wm = WorldModel(obs_space,step,conf,self._device)
        # elif self.wm_type=='v3':
            # self.wm = WorldModel_v3(obs_space,step,conf,self._device)

        # Actor critic

        # if self.wm_type=='v2':
        self.ac = ActorCritic(conf,self.wm,self._device)
        # elif self.wm_type=='v3':
        #     self.ac=ActorCritic_v3(conf,self.wm,self._device)
            #  self.ac = ActorCritic_v2(conf
            #                   )

        # Map probe

        # if conf.probe_model == 'map':
        #     probe_model = MapProbeHead(features_dim + 4, conf)
        # elif conf.probe_model == 'goals':
        #     probe_model = GoalsProbe(features_dim, conf)
        # elif conf.probe_model == 'map+goals':
        #     probe_model = MapGoalsProbe(features_dim, conf)
        # elif conf.probe_model == 'none':
        probe_model = NoProbeHead()
        # else:
        #     raise NotImplementedError(f'Unknown probe_model={conf.probe_model}')
        self.probe_model = probe_model
        self.probe_gradients = conf.probe_gradients

    def init_optimizers(self, lr, lr_actor=None, lr_critic=None, eps=1e-5,ac_eps=1e-5):
        if not self.probe_gradients:
            optimizer_wm = torch.optim.AdamW(self.wm.parameters(), lr=lr, eps=eps)
            optimizer_probe = torch.optim.AdamW(self.probe_model.parameters(), lr=lr, eps=eps)
            optimizer_actor = torch.optim.AdamW(self.ac.actor.parameters(), lr=lr_actor or lr, eps=ac_eps)
            optimizer_critic = torch.optim.AdamW(self.ac.critic.parameters(), lr=lr_critic or lr, eps=ac_eps)
            return optimizer_wm, optimizer_probe, optimizer_actor, optimizer_critic
        else:
            optimizer_wmprobe = torch.optim.AdamW(self.wm.parameters(), lr=lr, eps=eps)
            optimizer_actor = torch.optim.AdamW(self.ac.actor.parameters(), lr=lr_actor or lr, eps=ac_eps)
            optimizer_critic = torch.optim.AdamW(self.ac.critic.parameters(), lr=lr_critic or lr, eps=ac_eps)
            return optimizer_wmprobe, optimizer_actor, optimizer_critic

    def grad_clip(self, grad_clip, grad_clip_ac=None):
        if not self.probe_gradients:
            grad_metrics = {
                'grad_norm': nn.utils.clip_grad_norm_(self.wm.parameters(), grad_clip),
                'grad_norm_probe': nn.utils.clip_grad_norm_(self.probe_model.parameters(), grad_clip),
                'grad_norm_actor': nn.utils.clip_grad_norm_(self.ac.actor.parameters(), grad_clip_ac or grad_clip),
                'grad_norm_critic': nn.utils.clip_grad_norm_(self.ac.critic.parameters(), grad_clip_ac or grad_clip),
            }
        else:
            grad_metrics = {
                'grad_norm': nn.utils.clip_grad_norm_(self.wm.parameters(), grad_clip),
                'grad_norm_actor': nn.utils.clip_grad_norm_(self.ac.actor.parameters(), grad_clip_ac or grad_clip),
                'grad_norm_critic': nn.utils.clip_grad_norm_(self.ac.critic.parameters(), grad_clip_ac or grad_clip),
            }
        return grad_metrics

    def init_state(self, batch_size: int):
        return self.wm.init_state(batch_size)

    def inference(self,
                  obs: Dict[str, Tensor],
                  in_state: Any,
                  ) -> Tuple[D.Distribution, Any, Dict]:
        assert 'action' in obs, 'Observation should contain previous action'
        act_shape = obs['action'].shape
        assert len(act_shape) == 3 and act_shape[0] == 1, f'Expected shape (1,B,A), got {act_shape}'

        # Forward (world model)

        features, out_state = self.wm.forward(obs, in_state)

        # Forward (actor critic)
        if self.wm_type=='v2':

            feature = features[:, :, 0]  # (T=1,B,I=1,F) => (1,B,F) ## features.shape(1,1,1,2048)T*B*I*(D+S)
            action_distr=self.ac.forward_actor(feature)
            value=self.ac.forward_value(feature)
            metrics = dict(policy_value=value.detach().mean())
        elif self.wm_type=='v3':
            feature = features[:, :, 0]
            # feature = features[0, :, :] ##features.shape(1,1,1536)T*B*(D+S)
            action_distr = self.ac.forward_actor(feature)  # (1,B,A)
            value = self.ac.critic(feature)  # (1,B)
            metrics={}
            metrics.update(tensorstats(value.mode().detach(), "policy_value"))
            # metrics = dict(policy_value=value.detach().mean())
            # feature = features[:, :, 0]  # (T=1,B,I=1,F) => (1,B,F) ## features.shape(1,1,1,2048)T*B*I*(D+S)
            # action_distr=self.ac.forward_actor(feature)
            # value=self.ac.forward_value(feature)
            # metrics = dict(policy_value=value.detach().mean())

        return action_distr, out_state, metrics

    def training_step(self,
                      obs: Dict[str, Tensor],
                      in_state: Any,
                      iwae_samples: Optional[int] = None,
                      imag_horizon: Optional[int] = None,
                      do_open_loop=False,
                      do_image_pred=False,
                      do_dream_tensors=False,
                      ):
        assert 'action' in obs, '`action` required in observation'
        assert 'reward' in obs, '`reward` required in observation'
        assert 'reset' in obs, '`reset` required in observation'
        assert 'terminal' in obs, '`terminal` required in observation'
        iwae_samples = int(iwae_samples or self.iwae_samples)
        imag_horizon = int(imag_horizon or self.imag_horizon)
        T, B = obs['action'].shape[:2]
        I, H = iwae_samples, imag_horizon
        # print(T,B,I,H)
        # World model

        if self.wm_type=='v2':
            loss_model, features, states, out_state, metrics, tensors = \
                self.wm.training_step(obs,
                                    in_state,
                                    iwae_samples=iwae_samples,
                                    do_open_loop=do_open_loop,
                                    do_image_pred=do_image_pred)
        elif self.wm_type=='v3':
            # loss_model,features, post, context, metrics,tensors=self.wm.training_step(obs,do_image_pred=do_image_pred)
            # out_state={key: tensor[-1] for key, tensor in post.items()}
            loss_model, features, states, out_state, metrics, tensors = \
                self.wm.training_step(obs,
                                    in_state,
                                    iwae_samples=iwae_samples,
                                    do_open_loop=do_open_loop,
                                    do_image_pred=do_image_pred)

        # Map probe

        features_probe = features.detach() if not self.probe_gradients else features
        loss_probe, metrics_probe, tensors_probe = self.probe_model.training_step(features_probe, obs)
        metrics.update(**metrics_probe)
        tensors.update(**tensors_probe)

        # Policy

        in_state_dream: StateB = map_structure(states, lambda x: flatten_batch(x.detach())[0])
        # type: ignore  # (T,B,I) => (TBI)
        # if self.wm_type=='v2':
        features_dream, actions_dream, rewards_dream, terminals_dream, states_dream = \
        self.dream(in_state_dream, H, self.ac._conf.actor_grad == 'dynamics')  # (H+1,TBI,D)

        # elif self.wm_type=='v3':
        # #     # states=post
        # #     flatten = lambda x: x.reshape([-1] + list(x.shape[2:]))
        # #     in_state_dream = {k: flatten(v) for k, v in states.items()}
        # #     # in_state_dream={key: tensor[0] for key, tensor in states.items()}
        # #     # in_state_dream=map_structure(states, lambda x: flatten_batch(x.detach())[0])
        # # # Note features_dream includes the starting "real" features at features_dream[0]
        #     features_dream, actions_dream, rewards_dream, terminals_dream, states_dream = \
        #     self.dream(in_state_dream, H, self.ac._conf.actor_grad == 'dynamics')
            # states_dream={key: tensor.detach() for key, tensor in states_dream.items()}
        if self.wm_type=="v2":
            (loss_actor, loss_critic), metrics_ac, tensors_ac = \
                self.ac.training_step(features_dream.detach(),
                                    actions_dream.detach(),
                                    rewards_dream.mean.detach(),
                                    terminals_dream.mean.detach())
            # tensors.update(policy_value=unflatten_batch(tensors_ac['value'][0], (T, B, I)).mean(-1))
        elif self.wm_type=="v3":
            #  (loss_actor, loss_critic), metrics_ac, tensors_ac = \
            #     self.ac.training_step(features_dream.detach(),
            #                         actions_dream.detach(),
            #                         rewards_dream.mean.detach(),
            #                         terminals_dream.mean.detach(),
            #                         # states_dream
            #                         )
            (loss_actor, loss_critic), metrics_ac, tensors_ac = \
                self.ac.training_step(features_dream.detach(),
                                    actions_dream.detach(),
                                    rewards_dream.mode().detach(),
                                    terminals_dream.mode().detach())
            # tensors.update(policy_value=unflatten_batch(tensors_ac['value'][0], (T, B, I)).mean(-1))
        metrics.update(**metrics_ac)

        # Dream for a log sample.
        ## 扰动和刻意的action应该都可以在这里加。
        dream_tensors = {}
        if self.wm=="v2":
            if do_dream_tensors and self.wm.decoder.image is not None:
                with torch.no_grad():  # careful not to invoke modules first time under no_grad (https://github.com/pytorch/pytorch/issues/60164)
                    # The reason we don't just take real features_dream is because it's really big (H*T*B*I),
                    # and here for inspection purposes we only dream from first step, so it's (H*B).
                    # Oh, and we set here H=T-1, so we get (T,B), and the dreamed experience aligns with actual.
                    # 这里实际做的时候，T=1，只从第一步想象
                    in_state_dream: StateB = map_structure(states, lambda x: x.detach()[0, :, 0])  # type: ignore  # (T,B,I) => (B)
                    ## 基本上只改了这一步
                    # non_zero_indices = torch.nonzero(in_state_dream[1])

                    # print(non_zero_indices)
                    features_dream, actions_dream, rewards_dream, terminals_dream,states_dream = self.dream(in_state_dream, T - 1)  # H = T-1
                    # features_dream, actions_dream, rewards_dream, terminals_dream = self.dream_cond_action(in_state_dream, obs['action'])
                    image_dream = self.wm.decoder.image.forward(features_dream)
                    ## 拿Dreamer_agent的数据只训练actor_critic
                    # _, _, tensors_ac = self.ac.training_step(features_dream, actions_dream[1:,:,:], rewards_dream.mean, terminals_dream.mean, log_only=True)
                    _, _, tensors_ac = self.ac.training_step(features_dream, actions_dream, rewards_dream.mean, terminals_dream.mean, log_only=True)
                    # The tensors are intentionally named same as in tensors, so the logged npz looks the same for dreamed or not
                    # dream_tensors = dict(action_pred=actions_dream,  # first action is real from previous step
                    #                      reward_pred=rewards_dream.mean,
                    #                      terminal_pred=terminals_dream.mean,
                    #                      image_pred=image_dream,
                    #                      **tensors_ac)
                    dream_tensors = dict(action_pred=torch.cat([obs['action'][:1], actions_dream]),  # first action is real from previous step
                                        reward_pred=rewards_dream.mean,
                                        terminal_pred=terminals_dream.mean,
                                        image_pred=image_dream,
                                        **tensors_ac)
                    assert dream_tensors['action_pred'].shape == obs['action'].shape
                    assert dream_tensors['image_pred'].shape == obs['image'].shape
        elif self.wm_type=="v3":
            if do_dream_tensors:
                with torch.no_grad():  # careful not to invoke modules first time under no_grad (https://github.com/pytorch/pytorch/issues/60164)
                    # The reason we don't just take real features_dream is because it's really big (H*T*B*I),
                    # and here for inspection purposes we only dream from first step, so it's (H*B).
                    # Oh, and we set here H=T-1, so we get (T,B), and the dreamed experience aligns with actual.
                    # 这里实际做的时候，T=1，只从第一步想象
                    # in_state_dream={key: tensor[0] for key, tensor in states.items()}
                    in_state_dream: StateB = map_structure(states, lambda x: x.detach()[0, :, 0])
                    ## 基本上只改了这一步
                    # non_zero_indices = torch.nonzero(in_state_dream[1])

                    # print(non_zero_indices)
                    features_dream, actions_dream, rewards_dream, terminals_dream,states_dream = self.dream(in_state_dream, T - 1)  # H = T-1
                    # features_dream, actions_dream, rewards_dream, terminals_dream = self.dream_cond_action(in_state_dream, obs['action'])
                    image_dream = self.wm.decoder.image.forward(features_dream)
                    # image_dream = self.wm.decoder.image.forward(features_dream)
                    _, _, tensors_ac = self.ac.training_step(features_dream, actions_dream, rewards_dream.mode(), terminals_dream.mode(), log_only=True)
                    ## 拿Dreamer_agent的数据只训练actor_critic
                    # _, _, tensors_ac = self.ac.training_step(features_dream, actions_dream[1:,:,:], rewards_dream.mean, terminals_dream.mean, log_only=True)
                    # _, _, tensors_ac = self.ac.training_step(features_dream, actions_dream, rewards_dream.mean, terminals_dream.mean, log_only=True)
                    # (actor_loss,value_loss),features, states, actions, metrics=self.ac.training_step(features_dream.detach(),
                    #                 actions_dream.detach(),
                    #                 rewards_dream.detach(),
                    #                 terminals_dream.detach(),
                    #                 states_dream)
                    # The tensors are intentionally named same as in tensors, so the logged npz looks the same for dreamed or not
                    # dream_tensors = dict(action_pred=actions_dream,  # first action is real from previous step
                    #                      reward_pred=rewards_dream.mean,
                    #                      terminal_pred=terminals_dream.mean,
                    #                      image_pred=image_dream,
                    #                      **tensors_ac)
                    dream_tensors = dict(action_pred=torch.cat([obs['action'][:1], actions_dream]),  # first action is real from previous step
                                        reward_pred=rewards_dream.mode(),
                                        terminal_pred=terminals_dream.mode(),
                                        image_pred=image_dream,
                                        )
                    assert dream_tensors['action_pred'].shape == obs['action'].shape
                    assert dream_tensors['image_pred'].shape == obs['image'].shape
        if not self.probe_gradients:
            losses = (loss_model, loss_probe, loss_actor, loss_critic)
        else:
            losses = (loss_model + loss_probe, loss_actor, loss_critic)
        return losses, out_state, metrics, tensors, dream_tensors

    def dream_cond_action(self, in_state, actions, dynamics_gradients=False):
        features = []
        imag_horizon=actions.shape[0]-1
        state = in_state
        self.wm.requires_grad_(False)  # Prevent dynamics gradiens from affecting world model

        ## 如果想设计一些简单的任务，比如说一直往左转，那么可以在这里修改action的选取
        for i in range(imag_horizon):
            feature = self.wm.dynamics.to_feature(*state)
            # print(feature.shape)
            # action_dist = self.ac.forward_actor(feature)
            # if dynamics_gradients:
            #     action = action_dist.rsample()
            # else:
            #     action = action_dist.sample()
            features.append(feature)
            action=actions[i+1,:,:]
            # When using dynamics gradients, this causes gradients in RSSM, which we don't want.
            # This is handled in backprop - the optimizer_model will ignore gradients from loss_actor.
            _, state = self.wm.dynamics.cell.forward_prior(action, None, state)

        feature = self.wm.dynamics.to_feature(*state)
        features.append(feature)
        features = torch.stack(features)  # (H+1,TBI,D)
        # actions = torch.stack(actions)  # (H,TBI,A)

        rewards = self.wm.decoder.reward.forward(features)      # (H+1,TBI)
        terminals = self.wm.decoder.terminal.forward(features)  # (H+1,TBI)

        self.wm.requires_grad_(True)
        return features, actions, rewards, terminals

    def dream(self, in_state,imag_horizon: int, dynamics_gradients=False,perturb='None'):
        features = []
        actions = []
        states=[]
        state = in_state
        self.wm.requires_grad_(False)  # Prevent dynamics gradiens from affecting world model

        ## 如果想设计一些简单的任务，比如说一直往左转，那么可以在这里修改action的选取
        for i in range(imag_horizon):
            feature = self.wm.dynamics.to_feature(*state)
            if self.wm_type=='v2':
                action_dist = self.ac.forward_actor(feature)
            elif self.wm_type=="v3":
                # feature = self.wm.dynamics.to_feature(state)
                # feature=torch.cat((state[0], state[1]), -1)
                action_dist = self.ac.forward_actor(feature)
            if perturb=='guassian':
                # Get the batch shape and event shape of the given distribution.
                batch_shape = action_dist.batch_shape
                event_shape = action_dist.event_shape

                # Create a new standard normal distribution with the same structure.
                perturb_dist = torch.distributions.Normal(loc=torch.zeros(*batch_shape, *event_shape),
                                                    scale=torch.ones(*batch_shape, *event_shape))
                # perturb_dist =perturb_dist.to(action_dist.device)
                action=action_dist.sample()
                action=action+perturb_dist.sample().to(action.device)
            ## offline的情况
            elif perturb=='random_policy':
                # 需要生成的one-hot向量的形状
                action_sample=action_dist.sample()
                tensor_shape = action_sample.shape
                act_dim=action_sample.shape[1]

                # 生成一个介于[0, 18)的随机整数tensor，形状为(480, 4)
                random_indices = torch.randint(low=0, high=act_dim, size=(tensor_shape[0],))

                # 使用one_hot函数将随机整数tensor转换为one-hot编码
                action = F.one_hot(random_indices, num_classes=act_dim).to(torch.float32).to(action_sample.device)

                # 验证one_hot_tensor的形状

            ## '3' means be attacked
            elif perturb == 'attack_s':
                # Apply adversarial attack to the feature
                adversarial_feature = self.ac.adversarial_attack(feature,epsilon=0.2)

                # Compute action distribution for the adversarial feature
                action_dist_adv = self.ac.forward_actor(adversarial_feature)

                # Sample an action from the adversarial action distribution
                action = action_dist_adv.sample()
            elif perturb == 'attack_m':
                # Apply adversarial attack to the feature
                adversarial_feature = self.ac.adversarial_attack(feature,epsilon=0.8)

                # Compute action distribution for the adversarial feature
                action_dist_adv = self.ac.forward_actor(adversarial_feature)

                # Sample an action from the adversarial action distribution
                action = action_dist_adv.sample()
            elif perturb == 'attack_l':
                # Apply adversarial attack to the feature
                adversarial_feature = self.ac.adversarial_attack(feature,epsilon=2)

                # Compute action distribution for the adversarial feature
                action_dist_adv = self.ac.forward_actor(adversarial_feature)

                # Sample an action from the adversarial action distribution
                action = action_dist_adv.sample()
            elif perturb=='None':
                if dynamics_gradients:
                    action = action_dist.rsample()
                else:
                    action = action_dist.sample()
            features.append(feature)
            actions.append(action)
            states.append(state)
            # When using dynamics gradients, this causes gradients in RSSM, which we don't want.
            # This is handled in backprop - the optimizer_model will ignore gradients from loss_actor.
            # if self.wm_type=='v2':
            _, state = self.wm.dynamics.cell.forward_prior(action, None, state)
            # elif self.wm_type=='v3':
            #     state=self.wm.dynamics.img_step(state,action)

        # if self.wm_type=='v2':
            feature = self.wm.dynamics.to_feature(*state)
                # action_dist = self.ac.forward_actor(feature)
        # elif self.wm_type=="v3":
        #         feature = self.wm.dynamics.to_feature(state)
        #         # feature=torch.cat((state[0], state[1]), -1)
        #         # action_dist = self.ac.actor(feature)
        features.append(feature)
        states.append(state)
        features = torch.stack(features)  # (H+1,TBI,D)
        actions = torch.stack(actions)  # (H,TBI,A)
        # states=torch.stack(states)
        if self.wm_type=='v2':
            ##返回的是一个分布
            rewards = self.wm.decoder.reward.forward(features)      # (H+1,TBI)
            terminals = self.wm.decoder.terminal.forward(features)  # (H+1,TBI)
        elif self.wm_type=='v3':
            # rewards=self.wm.heads["reward"](features)
            # terminals=self.wm.heads["terminal"](features)
            rewards = self.wm.decoder.reward.forward(features)      # (H+1,TBI)
            terminals = self.wm.decoder.terminal.forward(features)  # (H+1,TBI)
            # # 获取字典中的键
            # keys =states[0].keys()

            # # 为每个键堆叠张量
            # states = {key: torch.stack([d[key] for d in states]) for key in keys}

        self.wm.requires_grad_(True)
        return features, actions, rewards, terminals,states

    def __str__(self):
        # Short representation
        s = []
        s.append(f'Model: {param_count(self)} parameters')
        if self.wm_type=='v2':
            for submodel in (self.wm.encoder, self.wm.decoder, self.wm.dynamics, self.ac, self.probe_model):
                if submodel is not None:
                    s.append(f'  {type(submodel).__name__:<15}: {param_count(submodel)} parameters')
        elif self.wm_type=="v3":
            for submodel in (self.wm.encoder, self.wm.decoder, self.wm.dynamics,self.ac):
                if submodel is not None:
                    s.append(f'  {type(submodel).__name__:<15}: {param_count(submodel)} parameters')
        return '\n'.join(s)

    def __repr__(self):
        # Long representation
        return super().__repr__()


# _*_ coding : utf-8 _*_
# @Time : 2023/7/18 17:16
# @Author : weixing
# @FileName : pytorch-minist
# @Project : cv

'''
准备数据集：使用 MNIST 数据集作为手写数字数据集。该数据集包含 60000 张 28x28 的灰度图像作为训练数据和 10000 张图像作为测试数据。
    将数据集分为训练集和测试集，比例为 6:1。
定义模型：使用 PyTorch 的 nn.Module 类定义一个神经网络模型。
编译模型：使用 PyTorch 的 nn.CrossEntropyLoss() 函数定义损失函数，使用 optimizer 类定义优化器，如 SGD 或 Adam。
        将模型、损失函数和优化器组合在一起，使用 model.train() 和 model.eval() 方法来训练和评估模型。
测试模型：使用测试集对模型进行评估。在每个 epoch 后，使用测试集计算模型的准确率和损失函数值，并记录最佳性能指标（如准确率或损失函数值）。
预测结果：使用训练好的模型对新的手写数字图像进行预测。将图像转换为合适的张量格式，输入模型中得到预测结果，并将预测结果转换为数字输出。
'''

import torch
import torch.nn as nn
from torchvision import datasets, transforms

# 准备数据集
def minist_loadData():
    train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transforms.ToTensor())
    test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transforms.ToTensor())
    print("训练数据集数量：", len(train_dataset))
    print("测试数据集数量：", len(test_dataset))
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)
    test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=64, shuffle=True)
    return train_loader, test_loader

"""
比较两个向量中有多少数字相等
"""
def count_equal_elements(vec1, vec2):
    count = 0
    for i in range(len(vec1)):
        if vec1[i] == vec2[i]:
            count += 1
    return count

class NetWork(nn.Module):
    def __init__(self):
        super(NetWork, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = nn.functional.relu(nn.functional.max_pool2d(self.conv1(x), 2))
        x = nn.functional.relu(nn.functional.max_pool2d(self.conv2(x), 2))
        x = x.view(-1, 320)
        x = nn.functional.relu(self.fc1(x))
        x = self.fc2(x)
        return nn.functional.log_softmax(x, dim=1)  #log_softmax下降效果显著，但是速度很慢；softmax下降效果不理想

class Model():
    def __init__(self, network):
        self.network = network
        # 定义损失函数、优化器
        self.criterion = nn.CrossEntropyLoss()
        self.optimizer = torch.optim.SGD(self.network.parameters(), lr=0.01)

    def train(self, epoch, train_loader):
        print("start training")
        # 训练模型
        for epoch in range(epoch):
            running_loss = 0.0
            for i, (inputs, labels) in enumerate(train_loader):
                self.optimizer.zero_grad()
                outputs = self.network(inputs)  # inputs [64, 1, 28, 28]  labels [64, 1]  outputs #[64, 10]
                loss = self.criterion(outputs, labels)
                loss.backward()
                self.optimizer.step()
                running_loss += loss.item()
            print('train;[Epoch %d] Loss: %.4f' % (epoch + 1, running_loss / len(train_loader)))

        print("train end")
    def eval(self, test_loader):
        print("start test")
        # 测试模型
        correct = 0
        total = 0
        with torch.no_grad():
            for inputs, labels in test_loader:
                outputs = self.network(inputs)  # [64, 10]
                _, predicted = torch.max(outputs.data, 1)  # [64]
                correct += count_equal_elements(labels, predicted)
                total += labels.size(0)
            print("predict;测试集总个数 %d, 预测正确的个数 %d, 准确率 %.4f" % (total, correct, correct / total))

            # 保存参数
            if correct/total > 0.97:
                print("model to be saved")
                torch.save(self.network.state_dict(), "./model/torch-mnist.h5")

        print("test end")
def main():
    # 加载数据集
    train_loader, test_loader = minist_loadData()
    # 网络
    network = NetWork()
    # 模型
    model = Model(network)
    # 训练
    epoch = 10
    model.train(epoch, train_loader)
    # 测试
    model.eval(test_loader)

if __name__=='__main__':
    main()


from abc import ABCMeta, abstractmethod

import six

@six.add_metaclass(ABCMeta)
class OrderedBinaryTree(object):
    @property
    @abstractmethod
    def root(self):
        pass

    @abstractmethod
    def delete(self, key):
        pass

    def update(self, key, val):
        node = self.search(key)
        if node is not None:
            node.val = val
        return node

    def search(self, key):
        return self._search(self._root, key)

    def _search(self, node, key):
        if node is None:
            return None
        if self._key_eq(node.key, key):
            return node
        if self._key_gt(node.key, key):
            return self._search(node.left, key)
        else:
            return self._search(node.right, key)

    def _right_rotate(self, node):
        l = node.left
        node.left = l.right
        if l.right is not None:
            l.right.parent = node
        l.right = node
        if not node.parent:
            self._root = l
        elif node is node.parent.left:
            node.parent.left = l
        else:
            node.parent.right = l
        l.parent = node.parent
        node.parent = l

    def _left_rotate(self, node):
        r = node.right
        node.right = r.left
        if r.left is not None:
            r.left.parent = node
        r.left = node
        if not node.parent:
            self._root = r
        elif node is node.parent.left:
            node.parent.left = r
        else:
            node.parent.right = r
        r.parent = node.parent
        node.parent = r


from datetime import datetime, timedelta
import json
import logging
import requests
import magic
from tenacity import retry, stop_after_attempt, wait_exponential
from urllib.parse import urlunsplit

logger = logging.getLogger(__file__)

def get_api_url(site_config, url):
    print("GET_API_URL")
    print(url)
    if url.find("www.eea.europa.eu") > -1:
#        if url.find("/api/") > -1 and not url.find("www.eea.europa.eu/en") > -1:
        if url.find("/api/") > -1:
            return url

    if site_config.get("fix_items_url", None):
        if f'{site_config["fix_items_url"]["without_api"]}/' in url:
            url = url.replace(
                site_config["fix_items_url"]["without_api"],
                site_config["fix_items_url"]["with_api"],
            )
        if url == site_config["fix_items_url"]["without_api"]:
            url = site_config["fix_items_url"]["with_api"]
        return url

    if site_config["url_api_part"].strip("/") == "":
        return url

    if f'/{site_config["url_api_part"]}' in url:
        logger.info(
            "Found url_api_part (%s) in url: %s",
            site_config["url_api_part"],
            url,
        )
        return url

    url_parts = url.split("/")
    if "://" in url:
        url_parts.insert(3, site_config["url_api_part"])
    else:
        url_parts.insert(1, site_config["url_api_part"])

    print("/".join(url_parts))
    return "/".join(url_parts)

def get_no_api_url(site_config, url):
    print("GET_NO_API_URL")
    print(url)
    if site_config.get('url_to_parse', None):
        return site_config.get('url_to_parse', None)
    if site_config.get("fix_items_url", None):
        if f'{site_config["fix_items_url"]["without_api"]}/' in url:
            return url
        if f'{site_config["fix_items_url"]["with_api"]}/' in url:
            return url.replace(
                site_config["fix_items_url"]["with_api"],
                site_config["fix_items_url"]["without_api"],
            )
        with_api2 = site_config["fix_items_url"].get("with_api2", None)
        if with_api2 is not None and f"{with_api2}/" in url:
            return url.replace(
                site_config["fix_items_url"]["with_api2"],
                site_config["fix_items_url"]["without_api"],
            )
        if url in [site_config["fix_items_url"].get("without_api",""), site_config["fix_items_url"].get("with_api",""), site_config["fix_items_url"].get("with_api2", "")]:
            return site_config["fix_items_url"]["without_api"]

    url_parts = url.split("://")
    protocol = url_parts[0]
    url = url_parts[1]

    ret_url = "/".join(url.split("/" + site_config["url_api_part"] + "/"))

    # Handle languages
    if url.find("www.eea.europa.eu") > -1:
        if url.find("/api/") > -1:
            ret_url = "/".join(ret_url.split("/api/"))
    print(ret_url)
    return f"{protocol}://{ret_url}"

def build_queries_list(site_config, query_config):
    query_limit = ""
    if query_config.get("quick", False):
        today = datetime.now()
        yesterday = today - timedelta(1)
        query_limit = f"&modified.query:date={yesterday.strftime('%m-%d-%Y')}&modified.range=min"
    url = site_config["url"].strip("/")
    if site_config.get("fix_items_url", None):
        if site_config["fix_items_url"]["without_api"] in url:
            url = url.replace(
                site_config["fix_items_url"]["without_api"],
                site_config["fix_items_url"]["with_api"],
            )
    else:
        url_api_part = site_config["url_api_part"].strip("/")
        if url_api_part != "":
            url = f"{url}/{url_api_part}"
    ts = datetime.now().timestamp()
    if site_config.get("portal_types", None):
        # queries = []
        queries = [
            f"{url}/@search?b_size={query_config['query_size']}&metadata_fields=modification_date&metadata_fields=modified&show_inactive=true&sort_order=reverse&sort_on=Date&portal_type={portal_type}{query_limit}&ts={ts}"
            for portal_type in site_config["portal_types"]
        ]
        if site_config.get("languages", None):
            for language in site_config.get("languages"):
                queries.append(
                    f"{url}/{language}/@search?b_size={query_config['query_size']}&metadata_fields=modification_date&metadata_fields=modified&show_inactive=true&sort_order=reverse&sort_on=Date{query_limit}&ts={ts}"
                )
    else:
        queries = [
            f"{url}/@search?b_size={query_config['query_size']}&metadata_fields=modification_date&metadata_fields=modified&show_inactive=true&sort_order=reverse&sort_on=Date{query_limit}&ts={ts}"
        ]
    return queries

@retry(wait=wait_exponential(), stop=stop_after_attempt(3))
def request_with_retry(url, method="get", data=None):
    print("Query:")
    print(url)
    print("-------------")
    logger.info("Fetching %s", url)
    handler = getattr(requests, method)
    resp = ""
    try:
        resp = handler(
            url,
            headers={"Accept": "application/json"},
            data=json.dumps(data),
            timeout=120,
        )
        logger.info("Response: %s", resp.text)
    except:
        logger.info("Timeout")

    assert json.loads(resp.text)  # test if response is json
    logger.info("Response is valid json")

    return resp.text

def execute_query(query):
    while True:
        resp = request_with_retry(query)
        docs = json.loads(resp)
        for doc in docs["items"]:
            yield (doc)
        next_query = docs.get("batching", {}).get("next", False)
        if next_query:
            query_base = query.split("@search")[0].split("#")[0]
            next_query_base = next_query.split("@search")[0].split("#")[0]
            query = next_query.replace(next_query_base, query_base)
        else:
            break

def get_docs(query):
    docs = execute_query(query)
    for doc in docs:
        yield (doc)

def get_doc_from_plone(site_config, doc_id):
    url_with_api = get_api_url(site_config, doc_id)
    r_url = f"{url_with_api}?expand=object_provides&eea_index=1"
    if site_config.get("avoid_cache_api", False):
        dt = datetime.now()
        # dts = datetime.strptime(
        #         dt.split("T")[0], "%Y-%m-%d"
        #     )
        r_url = f"{r_url}&crawler={dt}"

    r = request_with_retry(r_url)

    return r

@retry(wait=wait_exponential(), stop=stop_after_attempt(5))
def scrape_with_retry(v, url, js=False):
    logger.info("Scraping url: %s", url)
    hc = v.get("headless_chrome").get("endpoint")
    if js:
        resp = requests.post(
            hc,
            headers={"Content-Type": "application/json"},
            data=f'{{"url":"{url}", "js":true,"raw":true}}',
        )
        downloaded = resp.text
        status = resp.status_code
        final_url = resp.headers.get("final_url", "")
    else:
        resp = requests.get(url)
        downloaded = resp.text
        status = resp.status_code
        final_url = resp.headers.get("final_url", "")

    if magic.from_buffer(downloaded) == "data":
        return {
            "downloaded": None,
            "status_code": status,
            "final_url": final_url.split("?scrape")[0],
        }

    logger.info("Downloaded: %s", downloaded)

    return {
        "downloaded": downloaded,
        "status_code": status,
        "final_url": final_url.split("?scrape")[0],
    }

def scrape(v, site_config, doc_id):
    url_without_api = get_no_api_url(site_config, doc_id)
    scrape = False
    s_url = ""
    scrape_with_js = False
    if site_config.get("scrape_pages", False):
        s_url = url_without_api
        scrape_with_js = site_config.get("scrape_with_js", False)
        scrape = True
    response = {}
    if scrape:
        if site_config.get("avoid_cache_web", False):
            dt = datetime.now()
            # dts = datetime.strptime(
            #         dt.split("T")[0], "%Y-%m-%d"
            #     )
            s_url = f"{url_without_api}?scrape={dt}"
        response = scrape_with_retry(v, s_url, scrape_with_js)
    return response

FIELD_MARKERS = {"file": {"content-type", "download", "filename"}}

def is_field_of_type(field, _type):
    if not isinstance(field, dict):
        return False

    if _type not in FIELD_MARKERS:
        return False

    return set(field.keys()).issuperset(FIELD_MARKERS[_type])

def fix_download_url(download_url, source_url):
    if "www.eea.europa.eu" in source_url:
        return download_url.replace("@@download", "at_download")
    return download_url

CONTENT_TYPES_TO_EXTRACT = [
    "application/msword",
    "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
    "application/vnd.ms-word.document.macroEnabled.12",
]

def extract_attachments(json_doc, nlp_service_params, extract_pdf):

    url = json_doc["id"]
    params = nlp_service_params["converter"]

    logger.info("Extract attachments %s %s", json_doc, url)

    converter_dsn = urlunsplit(
        ("http", params["host"] + ":" + params["port"], params["path"], "", "")
    )

    text_fragments = []

    for name, value in json_doc.items():

        # if (
        #     is_field_of_type(value, "file")
        #     and value["content-type"] == "application/pdf"
        # ):
        extract_doc = False
        if is_field_of_type(value, "file"):
            if value["content-type"] == "application/pdf" and extract_pdf:
                extract_doc = True
            if value["content-type"] in CONTENT_TYPES_TO_EXTRACT:
                extract_doc = True

        if extract_doc:
            download_url = fix_download_url(value["download"], url)
            logger.info("Download url found: %s", download_url)
            try:
                resp = request_with_retry(
                    converter_dsn, "post", {"url": download_url}
                )
            except Exception:
                logger.exception("failed file extraction, retry")
                download_url = value["download"]
                logger.info("Retry with download url: %s", download_url)
                resp = request_with_retry(
                    converter_dsn, "post", {"url": download_url}
                )
            if isinstance(resp, str):
                resp = json.loads(resp)
            for doc in resp["documents"]:
                text_fragments.append(doc["text"].strip())

    text = "\n".join(text_fragments)

    logger.info("Retrieved file content: %r", text)

    return text

def extract_pdf(v, site_config, doc):
    pdf_text = ""
    nlp_service_params = v.get("nlp_services")
    should_extract_pdf = True

    if site_config.get("pdf_days_limit", 0) > 0:
        current_date = datetime.now()
        logger.info("CHECK DATE")
        mod_date_str = doc.get("modification_date", None)
        if mod_date_str:
            mod_date = datetime.strptime(
                mod_date_str.split("T")[0], "%Y-%m-%d"
            )
            logger.info(current_date)
            logger.info(mod_date)
            diff = current_date - mod_date
            delta = diff.days
            logger.info(delta)
            logger.info(site_config.get("pdf_days_limit"))
            if delta > site_config.get("pdf_days_limit"):
                should_extract_pdf = False
    # if should_extract_pdf:

    logger.info("EXTRACT DOCUMENT")
    logger.info(doc["id"])

    pdf_text = extract_attachments(doc, nlp_service_params, should_extract_pdf)

    return pdf_text


import time

import pytest
import sqlite3
from lyricfetch import Song

from conftest import Nothing

def test_config(database):
    assert database._connection
    table_query = 'select name from sqlite_master where type="table"'
    assert database._execute(table_query) == {'name': 'log'}

def test_execute(database):
    values = ('a', 'b', 'c')
    insert = 'insert into log (chat_id, artist, title) values (?, ?, ?)'
    ret = database._execute(insert, values)
    assert ret is None

    select = """
    select chat_id, artist, title from log
    where chat_id=? and artist=? and title=?
    """
    ret = database._execute(select, values)
    assert tuple(ret.values()) == values

def test_execute_connection_closed(database):
    """
    Test that the database will try to connect again if there are any errors.
    """
    values = ('a', 'b', 'c')
    insert = 'insert into log (chat_id, artist, title) values (?, ?, ?)'
    database._execute(insert, values)

    database.close()

    select = """
    select chat_id, artist, title from log
    where chat_id=? and artist=? and title=?
    """
    ret = database._execute(select, values)
    assert tuple(ret.values()) == values

def test_execute_connection_retries(database):
    """
    Test that the database will try to connect again if there are any errors.
    """
    select = 'select * from log'
    database._retries = 0
    with pytest.raises(sqlite3.Error):
        database._execute(select)

    database._retries = 1
    database.close()
    with pytest.raises(sqlite3.Error):
        database._execute(select)

def test_log_result(database):
    chat_id = 'chat_id'
    song = Song('lucis absentia', 'gehenna gate')
    source = Nothing()
    source.__name__ = 'source'
    result = Nothing()
    result.song = song
    result.source = source

    database.log_result(chat_id, result)
    query = database._execute('select chat_id, source, artist, title from log')
    assert query == dict(
        chat_id=chat_id, source='source', artist=song.artist, title=song.title
    )

    source.__name__ = 'new source'
    database.log_result(chat_id, result)
    query = database._execute('select chat_id, source, artist, title from log')
    assert query == dict(
        chat_id=chat_id,
        source='new source',
        artist=song.artist,
        title=song.title,
    )

def test_get_last_res(database):
    chat_id = 'id'
    artist = 'testament'
    title = 'a day of reckoning'
    source = 'source'
    album = 'the new order'
    now = int(time.time())
    insert = """
    insert into log
        (chat_id, artist, title, source, album, date)
        values (?, ?, ?, ?, ?, ?)
    """

    # Insert a record and retrieve it
    database._execute(insert, (chat_id, artist, title, source, album, now))
    database._connection.commit()
    assert database.get_last_res(chat_id) == dict(
        source=source, artist=artist, title=title, album=album
    )

    # Update the chat id with a newer entry
    artist = 'new artist'
    now += 1
    database._execute(insert, (chat_id, artist, title, source, album, now))
    assert database.get_last_res(chat_id) == dict(
        source=source, artist=artist, title=title, album=album
    )

    # Updating with an older entry should return the new one again
    artist = 'old artist'
    now -= 2
    database._execute(insert, (chat_id, artist, title, source, album, now))
    assert database.get_last_res(chat_id) == dict(
        source=source, artist='new artist', title=title, album=album
    )

@pytest.mark.parametrize('param, expect', [(1, '1'), ("'hello'", "''hello''")])
def test_sanitize(database, param, expect):
    assert database.sanitize(param) == expect


import bpy

bl_info = {
    "name": "Ngx Select",
    "author": "SolPie",
    "version": (1,0),
    "blender": (2, 68, 0),
    "location": "View3D > Specials > Ngx Select ",
    "description": "Select the verts on -x",
    "warning": "",
    "category": "Mesh"}

def sel_ng_x_verts(context):
    ob = context.active_object
    #sel_mode = bpy.context.tool_settings.mesh_select_mode
    bpy.ops.mesh.select_all(action='DESELECT')
    bpy.ops.object.mode_set(mode='OBJECT')
    #
    bpy.context.tool_settings.mesh_select_mode = [True, False, False]
    #num = 0
    for v in ob.data.vertices:
        if v.co.x < 0:
            #print('select -x:',v.co.x)
            v.select = True
            #num += 1
            #print(num)
    bpy.ops.object.mode_set(mode='EDIT')
    #bpy.context.tool_settings.mesh_select_mode = sel_mode

def set_mirror(v1, v2):
    bpy.ops.object.mode_set(mode='OBJECT')
    v2.co = v1.co
    v2.co.x = -v2.co.x
    bpy.ops.object.mode_set(mode='EDIT')

class NgxSel(bpy.types.Operator):
    bl_idname = 'mesh.ngx'
    bl_label = 'Ngx Select'
    bl_options = {'REGISTER', 'UNDO'}

    @classmethod
    def poll(cls, context):
        obj = context.active_object
        return (obj and obj.type == 'MESH')

    def execute(self, context):
        sel_ng_x_verts(context)
        return {'FINISHED'}

def menu_func(self, context):
    self.layout.operator(NgxSel.bl_idname, text="Ngx Select")

def register():
    bpy.utils.register_module(__name__)

    bpy.types.VIEW3D_MT_edit_mesh_specials.append(menu_func)
    bpy.types.VIEW3D_MT_edit_mesh_vertices.append(menu_func)

def unregister():
    bpy.utils.unregister_module(__name__)

    bpy.types.VIEW3D_MT_edit_mesh_specials.remove(menu_func)
    bpy.types.VIEW3D_MT_edit_mesh_vertices.remove(menu_func)

if __name__ == "__main__":
    register()


import os
import json
import binascii
import hashlib
from functools import wraps

from flask import current_app

class PollManager(object):
    """A very simple poll manager that saves its data as json."""
    def __init__(self, path):
        self.file = os.path.join(path, 'polls.json')

    def read(self):
        print("JSON path:")
        print(self.file)
        if not os.path.exists(self.file):
            return {}
        with open(self.file) as f:
            data = json.loads(f.read())
        return data

    def write(self, data):
        with open(self.file, 'w') as f:
            f.write(json.dumps(data, indent=2))

    def add_poll(self, name, title, options = [], votes = []):
        polls = self.read()
        if polls.get(name):
            return False
        new_poll = {
            'title': title,
            'options': options,
            'votes': votes
        }
        polls[name] = new_poll
        self.write(polls)
        polldata = polls.get(name)
        return Poll(self, name, polldata)

    def get_poll(self, name):
        polls = self.read()
        polldata = polls.get(name)
        if not polldata:
            return None
        return Poll(self, name, polldata)

    def delete_poll(self, name):
        polls = self.read()
        if not polls.pop(name, False):
            return False
        self.write(polls)
        return True

    def update(self, name, polldata):
        data = self.read()
        data[name] = polldata
        self.write(data)

    def get_basic_data(self, name):
        polls = self.read()
        polldata = polls.get(name)
        if not polldata:
            return None
        return Poll(self, name, polldata)

class Poll(object):
    def __init__(self, manager, name, data):
        self.manager = manager
        self.name = name
        self.data = data

    def get(self, option):
        return self.data.get(option)

    def set(self, option, value):
        self.data[option] = value
        self.save()

    def save(self):
        self.manager.update(self.name, self.data)

    def get_id(self):
        return self.name

class PollBasic(object):
    def __init__(self, name, title):
        self.name = name
        self.title = title


"""Methods for pre-processing subjects, e.g., pruning."""

import logging

from typing import Iterable, List, Mapping, Set

from slub_docsa.common.subject import SimpleSubjectHierarchy, SubjectHierarchy, SubjectTargets, SubjectTuple
from slub_docsa.common.subject import SubjectUriList

logger = logging.getLogger(__name__)

def subject_ancestors_list(
    subject_uri: str,
    subject_hierarchy: SubjectHierarchy,
) -> Iterable[str]:
    """Return the list of ancestors for a subject node in a subject hierarchy (including the subject itself).

    Parameters
    ----------
    subject_uri: str
        the uri of the subject whose ancestors are being returned
    subject_hierarchy: SubjectHierarchy
        the subject hierarchy which is being searched for ancestors of the specified subject

    Returns
    -------
    Iterable[str]
        the list of ancestors (starting with the root ancestors, ending with and including the subject itself)
    """
    ancestors: List[str] = []
    next_subject = subject_uri

    while next_subject is not None:
        ancestors.insert(0, next_subject)
        next_subject = subject_hierarchy.subject_parent(next_subject)

    return ancestors

def subject_siblings_list(
    subject_uri: str,
    subject_hierarchy: SubjectHierarchy,
) -> Iterable[str]:
    """Return the list of siblings for a subject node in a subject hierarchy.

    Parameters
    ----------
    subject_uri: str
        the uri of the subject whose siblings are being returned
    subject_hierarchy: SubjectHierarchy
        the subject hierarchy that is searched for siblings of the specified subject

    Returns
    -------
    Iterable[str]
        the list of siblings in arbitrary order, including the subject itself
    """
    siblings: List[str] = []
    subject_parent = subject_hierarchy.subject_parent(subject_uri)

    for siblings_uri in iter(subject_hierarchy):
        siblings_parent = subject_hierarchy.subject_parent(siblings_uri)

        if siblings_parent == subject_parent:
            siblings.append(siblings_uri)

    return siblings

def root_subjects_from_subject_parent_map(subject_parent: Mapping[str, str]) -> Iterable[str]:
    """Extract root subjects that do not have a parent.

    Parameters
    ----------
    subject_parent : Mapping[str, str]
        the mapping from subject URIs to parent subject URIs

    Returns
    -------
    Iterable[str]
        the list of root subject that do have `None` as parent URI in the provided mapping
    """
    root_subjects = set()
    for child_uri, parent_uri in subject_parent.items():
        if parent_uri is None:
            root_subjects.add(child_uri)
    return list(root_subjects)

def subject_label_breadcrumb(
    subject_uri: str,
    subject_hierarchy: SubjectHierarchy,
) -> Iterable[Mapping[str, str]]:
    """Return a list of label mappings describing the subject hierarchy path from root to this subject.

    Parameters
    ----------
    subject_uri: str
        the uri of the subject whose ancestor path (breadcrumb) is supposed to be returned
    subject_hierarchy: SubjectHierarchy
        the subject hierarchy that is searched for ancestors of the specified subject

    Returns
    -------
    Iterable[Mapping[str, str]]
        a simple breadcrumb string (e.g. subject1 | subject2 | subject3) describing the ancestor path of the subject
    """
    subject_path = subject_ancestors_list(subject_uri, subject_hierarchy)
    return list(map(subject_hierarchy.subject_labels, subject_path))

def subject_label_as_string(
    subject_uri: str,
    lang_code: str,
    subject_hierarchy: SubjectHierarchy,
) -> str:
    """Return a string label for a subject in the given language.

    Parameters
    ----------
    subject_uri: str
        the uri of the subject whose label is requested
    lang_code: str
        the ISO 639-1 language code of the language the label is requested in
    subject_hierarchy: SubjectHierarchy
        the subject hierarchy that contains the information about the subject

    Returns
    -------
    str
        a label for the subject in the given language or a short text indicating
        that there is no label available for the subject in the requested language
    """
    labels = subject_hierarchy.subject_labels(subject_uri)
    if labels is None:
        return f"no label for subject '{subject_uri}'"
    if lang_code not in labels:
        return f"label not available in language '{lang_code}' for subject '{subject_uri}'"
    return labels[lang_code]

def subject_label_breadcrumb_as_string(
    subject_uri: str,
    lang_code: str,
    subject_hierarchy: SubjectHierarchy,
) -> Iterable[str]:
    """Return a list of label mappings describing the subject hierarchy path from root to this subject.

    Parameters
    ----------
    subject_uri: str
        the uri of the subject whose ancestor path (breadcrumb) is supposed to be returned
    lang_code: str
        the ISO 639-1 language code of the language the breadcrumb is requested in
    subject_hierarchy: SubjectHierarchy
        the subject hierarchy that is searched for ancestors of the specified subject

    Returns
    -------
    Iterable[str]
        a simple breadcrumb string (e.g. ["subject1", "subject2", "subject3"]) describing the ancestor path of the
        subject including the requested subject itself
    """
    subject_path = subject_ancestors_list(subject_uri, subject_hierarchy)
    return list(subject_label_as_string(subject_uri, lang_code, subject_hierarchy) for subject_uri in subject_path)

def subject_ancestors_for_subject_list(
    subject_list: Iterable[str],
    subject_hierarchy: SubjectHierarchy,
) -> Iterable[str]:
    """Return the set of all ancestors of a list of subjects from a subject hierarchy.

    Can be used to reduce a large subject hierarchy to a smaller set of relevant subjects.

    Parameters
    ----------
    subject_list : Iterable[str]
        the list of subjects whose ancestors are determined
    subject_hierarchy : SubjectHierarchy
        the subject hierarchy that is used to find ancestors for a subject

    Returns
    -------
    Iterable[str]
        the set of subject URIs for the provided subjects and all of their ancestors
    """
    return set(
        a_uri for s_uri in subject_list for a_uri in subject_ancestors_list(s_uri, subject_hierarchy)
    )

def prune_subject_uri_to_level(
    level: int,
    subject_uri: str,
    subject_hierarchy: SubjectHierarchy
) -> str:
    """Return the ancestor subject uri at a specific hierarchy level.

    Parameters
    ----------
    level: int
        The level at which to return a ancestor subject uri. Level 1 is the root level.
    subject_uri: str
        The uri of the subject that is supposed to be pruned, i.e., replaced by a ancestor subject
    subject_hierarchy: SubjectHierarchy
        The subject hierarchy used for pruning

    Returns
    -------
    str
        Either the ancestor subject uri of an ancestor at the specified hierarchy level, or the original subject uri
        if the requested subject is already lower (near the root) in the hierarchy than the requested level.
    """
    if level < 1:
        raise ValueError("level must be 1 or larger than 1")
    if subject_uri not in subject_hierarchy:
        raise ValueError(f"subject uri {subject_uri} not in hiearchy")

    ancestors = list(subject_ancestors_list(subject_uri, subject_hierarchy))
    if len(ancestors) <= level:
        # subject is below (closer to root) or exactly at the requested level
        return subject_uri

    # subject must be above (closer to leaf) than level
    return ancestors[level - 1]

def prune_subject_uris_to_level(
    level: int,
    subjects: SubjectUriList,
    subject_hierarchy: SubjectHierarchy,
) -> SubjectUriList:
    """Prune a list of subjects to a specified hierarchy level.

    Duplicate subjects are removed. Applies `prune_subject_uri_to_level` to every subject uri of the provided list.

    Parameters
    ----------
    level: int
        The level at which to return a ancestor subject uri. Level 1 is the root level.
    subjects: SubjectUriList
        The list of subjects to be pruned to the specified hierarchy level
    subject_hierarchy: SubjectHierarchy
        The subject hierarchy that contains the required subject relationships

    Returns
    -------
    SubjectUriList
        a list of subject uris, each pruned according to `prune_subject_uri_to_level`
    """
    return list({prune_subject_uri_to_level(level, s, subject_hierarchy) for s in subjects})

def prune_subject_targets_to_level(
    level: int,
    subject_target_list: SubjectTargets,
    subject_hierarchy: SubjectHierarchy,
) -> SubjectTargets:
    """Prune all subjects of a subject list to the specified hierarchy level.

    Duplicate subjects are removed. Applies `prune_subject_uris_to_level` to every subject uri of the provided list of
    subject lists.

    Parameters
    ----------
    level: int
        The level at which to return a ancestor subject uri. Level 1 is the root level.
    subject_target_list: SubjectTargets
        The list of subject lists to be pruned to the specified hierarchy level
    subject_hierarchy: SubjectHierarchy
        The subject hierarchy that contains the required subject relationships

    Returns
    -------
    SubjectTargets
        A list of subject uri lists, each pruned according to `prune_subject_uris_to_level`
    """
    return list(map(
        lambda subjects: prune_subject_uris_to_level(level, subjects, subject_hierarchy),
        subject_target_list
    ))

def count_number_of_samples_by_subjects(subject_targets: SubjectTargets) -> Mapping[str, int]:
    """Count the number of occurances of subjects annotations.

    Parameters
    ----------
    subject_targets: SubjectTargets
        A list of subject lists as provided by a dataset

    Returns
    -------
    Mapping[str, int]
        Occurance counts for each subject
    """
    counts = {}

    for subject_list in subject_targets:
        for subject_uri in subject_list:
            counts[subject_uri] = counts.get(subject_uri, 0) + 1

    return counts

def children_map_from_subject_parent_map(
    subject_parent: Mapping[str, str],
) -> Mapping[str, Iterable[str]]:
    """Return a set of children subjects for each parent subject in a subject hierarchy.

    Parameters
    ----------
    subject_parent: Mapping[str, str]
        The mapping of a subject to its parent subject

    Returns
    -------
    Mapping[str, Iterable[str]]
        A mapping from parent subjects to their respective set of child subjects
    """
    subject_children: Mapping[str, Set[str]] = {}
    for subject_uri, parent_uri in subject_parent.items():
        if subject_uri not in subject_children:
            subject_children[subject_uri] = set()
        if parent_uri is not None:
            if parent_uri not in subject_children:
                subject_children[parent_uri] = set()
            subject_children[parent_uri].add(subject_uri)
    return subject_children

def build_subject_hierarchy_from_subject_tuples(
    subject_tuples: Iterable[SubjectTuple]
) -> SubjectHierarchy:
    """Build a subject hierarchy from a list of subject tuples.

    Parameters
    ----------
    subject_tuples : Iterable[SubjectTuple]
        the list of subject tuples containing information about each subjects URI, labels, parent subject

    Returns
    -------
    SubjectHierarchy
        the generated subject hierarchy
    """
    subject_parent: Mapping[str, str] = {}
    subject_labels: Mapping[str, Mapping[str, str]] = {}
    subject_notation: Mapping[str, str] = {}

    for subject_tuple in subject_tuples:
        subject_parent[subject_tuple.subject_uri] = subject_tuple.parent_uri
        subject_labels[subject_tuple.subject_uri] = subject_tuple.labels
        subject_notation[subject_tuple.subject_uri] = subject_tuple.notation

    root_subjects = root_subjects_from_subject_parent_map(subject_parent)
    subject_children = children_map_from_subject_parent_map(subject_parent)

    return SimpleSubjectHierarchy(root_subjects, subject_labels, subject_parent, subject_children, subject_notation)

def prune_subject_uris_to_parent(
    subject_uris: SubjectUriList,
    to_be_pruned_subjects: Set[str],
    subject_hierarchy: SubjectHierarchy,
) -> SubjectUriList:
    """Replace subject uris with uri of parent subject for those that are supposed to be pruned.

    Is used by `prune_subject_targets_to_minimum_samples` to incrementally replace subjects with their parent subjects
    in case there would be insufficient training examples for that subject.

    Parameters
    ----------
    subject_uris: SubjectUriList
        the list of subjects that is scanned for subjects that are supposed to be pruned
    to_be_pruned_subjects: Set[str]
        the list of subjects that is supposed to be replaced with their parent subjects
    subject_hierarchy: SubjectHierarchy
        The subject hierarchy that contains the required subject relationships

    Returns
    -------
    SubjectUriList
        the same subject list as `subject_uris`,  where some subjects are replaced with their parent subject
    """
    new_subject_uris = []
    for subject_uri in subject_uris:
        if subject_uri in to_be_pruned_subjects:
            parent_uri = subject_hierarchy.subject_parent(subject_uri)
            if parent_uri is not None:
                new_subject_uris.append(parent_uri)
            else:
                new_subject_uris.append(subject_uri)
        else:
            new_subject_uris.append(subject_uri)
    return list(set(new_subject_uris))

def prune_subject_targets_to_parent(
    subject_targets: SubjectTargets,
    to_be_pruned_subjects: Set[str],
    subject_hierarchy: SubjectHierarchy,
) -> SubjectTargets:
    """Replace subject uris in target list with uri of parent subject for those that are supposed to be pruned.

    Is used by `prune_subject_targets_to_minimum_samples` to incrementally replace subjects with their parent subjects
    in case there would be insufficient training examples for that subject.

    Parameters
    ----------
    subject_targets: SubjectTargets
        the list of subject lists that is scanned for subjects that are supposed to be pruned
    to_be_pruned_subjects: Set[str]
        the list of subjects that is supposed to be replaced with their parent subjects
    subject_hierarchy: SubjectHierarchy
        The subject hierarchy that contains the required subject relationships

    Returns
    -------
    SubjectTargets
        the same subject target list as `subject_targets`,  where some subjects are replaced with their parent subject
    """
    return list(
        map(lambda t: prune_subject_uris_to_parent(t, to_be_pruned_subjects, subject_hierarchy), subject_targets)
    )

def prune_subject_targets_to_minimum_samples(
    minimum_samples: int,
    subject_targets: SubjectTargets,
    subject_hierarchy: SubjectHierarchy
) -> SubjectTargets:
    """Prune subject targets such that subjects with insufficient samples are replaced with their parents.

    Subjects with no parents are left as is. Therefore, it is not guaranteed that all subjects will have a minimum
    number of samples after pruning.

    Parameters
    ----------
    minimum_samples: int
        minimum number of required samples for subjects to be considered of acceptable size
    subject_targets: SubjectTargets
        subject targets that are checked for their number of samples, and modified if necessary (not in-place)
    subject_hierarchy: SubjectHierarchy
        The subject hierarchy that contains the required subject relationships

    Returns
    -------
    SubjectTargets
        the same subject target list as `subject_targets`,  where some subjects are replaced with their parent subject
        in case they do not have a sufficient number of samples
    """
    def _relevant_subject_children(relevant_subjects, subject_hierarchy):
        subject_children = {
            s_uri: [c_uri for c_uri in subject_hierarchy.subject_children(s_uri) if c_uri in relevant_subjects]
            for s_uri in relevant_subjects
        }
        return {s_uri: children for s_uri, children in subject_children.items() if children}

    pruned_subject_targets = subject_targets
    subject_counts = count_number_of_samples_by_subjects(subject_targets)
    relevant_subjects = subject_ancestors_for_subject_list(list(subject_counts.keys()), subject_hierarchy)
    subject_children = _relevant_subject_children(relevant_subjects, subject_hierarchy)
    subjects_to_be_checked = set(iter(relevant_subjects))

    i = 0
    while len(subjects_to_be_checked) > 0:
        # only consider childless subjects, whose counts can not increase further
        childless_subjects = [s for s in subjects_to_be_checked if s not in subject_children]
        logger.debug("prune to parent iteration %d considering %d childless subjects", i + 1, len(childless_subjects))

        # find subjects that need pruning
        childless_subjects_to_be_pruned = {
            s for s in childless_subjects if s in subject_counts and subject_counts[s] < minimum_samples
        }

        # do pruning
        logger.debug("prune %d subjects because of insufficient samples", len(childless_subjects_to_be_pruned))
        pruned_subject_targets = prune_subject_targets_to_parent(
            pruned_subject_targets,
            childless_subjects_to_be_pruned,
            subject_hierarchy
        )

        # check each childless subject one by one
        for s_uri in childless_subjects:
            # remove every childless subject from child set of parent,
            # such that parent can become childless in the next iterations,
            # and eventually all subjects are checked
            parent_uri = subject_hierarchy.subject_parent(s_uri)
            if parent_uri:
                subject_children[parent_uri].remove(s_uri)
                if len(subject_children[parent_uri]) == 0:
                    del subject_children[parent_uri]

        # update counts and check again
        subject_counts = count_number_of_samples_by_subjects(pruned_subject_targets)
        subjects_to_be_checked.difference_update(childless_subjects)
        i += 1

    return pruned_subject_targets