#####################################################################################
#extract.py
#Authors: Revanth Pobala and Hussain Mucklai.
#This project is about the Testing of Lilac(Lightweight low latency anonymous chat)
#This program will provide the graph for cdf(Cumulative Distribution function)
#####################################################################################
import os
import matplotlib.pyplot as plt
from pylab import *
from scipy.stats import norm
import matplotlib.patches as mpatches
from bisect import bisect_left

class discrete_cdf:
    def __init__(self,data):
        self._data = data # must be sorted
        self._data_len = float(len(data))

    def __call__(self,point):
        return (len(self._data[:bisect_left(self._data, point)]) /
                self._data_len)

class extract():

    def __init__(self):
        self.to_draw = []

    def read_from_file(self,filename):
        file = open(filename,'rw')
        return file

    def merge_files(self):
        os.system('cat /home/revanth/Projects/workspace/Lilac_Testing/src/users/logs/*.txt >/home/revanth/Projects/workspace/Lilac_Testing/src/Results/finalfile.log')

    def add_to_dict(self):
        items = {}
        graph_values = []
        to_draw = []
        median = []
        sum = 0
        file = extract.read_from_file('finalfile.log')
        for i in file.readlines():
            if len(i.strip()) != 0:
                content = i.split(',')
                msg = content[0].strip()
                Time = float(str(content[1].strip()))
                if msg not in items:
                    items[msg] = []
                items[msg].append(Time)

        for i in items:
            for j in items[i]:
                graph_values.append(float(max(items[i])- min(items[i])))
        #graph_values = set(graph_values)
        for i in graph_values:
            if i >= 1.5:
                to_draw.append(i*0.1)
                sum = sum+i
        avg = np.average(sum)/len(to_draw)*0.1
        #plt.plot(0,avg)
        self.to_draw = np.array(to_draw)
        anthax =  np.arange(0,len(to_draw))
        plt.plot(anthax, self.to_draw, '-')
        plt.ylabel('Time (Seconds)')
        plt.xlabel('Number of Messages')
        plt.title(r'Time  vs Number of messages')
        plt.show()

    def draw_cdf(self):
        self.to_draw = sort(self.to_draw)
        #print self.to_draw
        cdf = discrete_cdf(self.to_draw)
        xvalues =  np.arange(0,6)
        yvalues = [cdf(point) for point in xvalues]
        plt.plot( yvalues)
        #plt.xticks(xvalues)
       # red_patch = mpatches.Patch(color='green', label='')
       # plt.legend(handles=[red_patch])
        plt.title(r'Cumulative Distribution Function')
        plt.ylabel('cdf of messages')
        plt.xlabel('Time')
        plt.show()

extract =  extract()
extract.merge_files()
extract.add_to_dict()
extract.draw_cdf()


from cafle import Account
from functools import partial

from astn0_index import idx
from astn1_account import acc

loan = Account(idx)
with loan as l:
    l.mtrt = idx.mtrt
    l.idxfn = idx.loan
    l.is_wtdrbl = False
    l.is_repaid = False

loan.ntnl = loan.subacc('ntnl')
with loan.ntnl as l:
    l.amt = 65_000
    l.subscd(loan.idxfn[0], l.amt)
    l.addscd(loan.idxfn[-1], l.amt)

loan.IR = loan.subacc('IR')
with loan.IR as l:
    l.rate = 0.08
    l.cycle = 1
    l.rate_cycle = l.rate / 12 * l.cycle

loan.fee = loan.subacc('fee')
with loan.fee as l:
    l.rate = 0.02

#Estimate Loan Cost
def estimate_fee_amt(loan, idxno):
    if idxno == loan.idxfn[0]:
        feeamt = loan.ntnl.amt * loan.fee.rate
        loan.fee.addscd(idxno, feeamt, note='fee_amt')
        return feeamt
    return 0
loan.estimate_fee_amt = partial(estimate_fee_amt, loan)

def estimate_IR_amt(loan, idxno):
    if loan.is_wtdrbl is False:
        return 0
    if loan.is_repaid is True:
        return 0
    ntnlbal = -loan.ntnl.bal_strt[idxno]
    IRamt = ntnlbal * loan.IR.rate_cycle
    if IRamt > 0.0:
        loan.IR.addscd(idxno, IRamt, note='IR_amt')
        return IRamt
    return 0
loan.estimate_IR_amt = partial(estimate_IR_amt, loan)

#Withdraw Loan Amount
def set_loan_withdrawable(loan, idxno):
    if idxno == loan.idxfn[0]:
        loan.is_wtdrbl = True
loan.set_loan_withdrawable = partial(set_loan_withdrawable, loan)

def withdraw_ntnl_fixed(loan, acc, idxno):
    amt_wtdrw = loan.ntnl.scd_out[idxno]
    loan.ntnl.send(idxno, amt_wtdrw, acc, "withdraw loan")
loan.withdraw_ntnl_fixed = partial(withdraw_ntnl_fixed, loan)

#Pay Loan Costs
def pay_fee_amt(loan, acc, idxno):
    feeamt = loan.fee.scd_in[idxno]
    acc.send(idxno, feeamt, loan.fee, f"pay fee amt at {idxno}")
    return feeamt
loan.pay_fee_amt = partial(pay_fee_amt, loan)

def pay_IR_amt(loan, acc, idxno):
    IRamt = loan.IR.scd_in[idxno]
    acc.send(idxno, IRamt, loan.IR, f"pay IR amt at {idxno}")
    return IRamt
loan.pay_IR_amt = partial(pay_IR_amt, loan)

#Repay Loan Amount
def repay_ntnl_amt(loan, acc, idxno):
    #at maturity
    if idxno >= idx.loan[-1]:
        amt_scd_in = loan.ntnl.rsdl_in_cum[idxno]
        acc.send(idxno, amt_scd_in, loan.ntnl, "repay loan")
loan.repay_ntnl_amt = partial(repay_ntnl_amt, loan)

def setback_loan_unwithdrawable(loan, idxno):
    if idxno >= idx.loan[-1]:
        loan.is_wtdrbl = False
loan.setback_loan_unwithdrawable = partial(setback_loan_unwithdrawable, loan)


# Temporary, for test purposes
import os
os.environ['APP_CONFIG'] = '/brum/dev/sharp_eye/src/resources/cam1.yaml'

from sharp_eye import RtspCamera
from time import sleep
from datetime import datetime
import numpy

from lib.quicklock import lock
lock()

DELAY = 1.0

cam = RtspCamera()

images = 0
times = []

def get_snapshot():
    snapshot_start = datetime.now()
    img = cam.snapshot_img(timeout=5)
    if (datetime.now() - snapshot_start).total_seconds() < DELAY:
        sleep(DELAY - (datetime.now() - snapshot_start).total_seconds())
    return img

get_snapshot()

for i in range(1, 999201):
    start = datetime.now()
    img = get_snapshot()
    time = (datetime.now() - start).total_seconds()
    times.append(time)
    if img is None:
        print('%s: %s No image ...' % (i, time))
        continue
    images += 1
    round_time = round(time, 2)
    print('%s: %s - %s : stats {avg: %.2f, delayed: %s, max_delay: %.2f avg_delay: %.2f} ' % (
        i,
        round_time,
        img.shape,
        numpy.mean(times),
        len([x for x in times if x > DELAY + 0.5]),
        max(times),
        numpy.mean([x for x in times if x > DELAY + 0.5])
    ))

print('Got %s images' % images)
print('Average image retrieval time - %s' % numpy.mean(times))


from django.urls import path, include, re_path

from rest_framework.routers import DefaultRouter

from .views import ProfileViewSet

router = DefaultRouter()

router.register(
    r'profile',
    ProfileViewSet,
    basename="profile"
)

urlpatterns = [
    re_path(r'', include(router.urls))
]


import os
from copy import deepcopy
from typing import List

from to_typescript.core.filter import FilterAbstract
from to_typescript.core.transform.extra_rules import ClassInheritance
from to_typescript.filters.function_save import FilterFunctionSave

class FilterOOPSave(FilterAbstract):
    FILE_STARTER = FilterFunctionSave.FILE_STARTER
    DUMP_FOLDERS = dict(server='output/server/oop',
                        client='output/client/oop',
                        client_gui='output/client/oop/gui')

    imports = FilterFunctionSave.imports

    @staticmethod
    def shift_lines(text: str, spaces: int) -> str:
        return '\n'.join(
            ' ' * spaces + line if line else line
            for line in text.split('\n')
        )

    def __init__(self):
        super().__init__()

        self.imports = deepcopy(self.imports)
        for key in self.imports['shared']:
            self.imports['client'].append(key)
            self.imports['server'].append(key)

    @staticmethod
    def get_class_begin(class_name: str, generic_string: str) -> str:
        inherit = ClassInheritance(class_name).get_child()
        if inherit:
            inherit_inline = f' extends {inherit}'
        else:
            inherit_inline = ''

        return f'''
/** @customConstructor {class_name} */
export class {class_name}{generic_string}{inherit_inline} {{
'''

    @staticmethod
    def get_class_end() -> str:
        return '}\n'

    def save_file_category(self,
                           class_name: str, side: str,
                           fields: List[str],
                           methods: List[str],
                           class_templates: List[str]):
        if not fields and not methods:
            return

        is_gui = class_name.startswith('Gui')
        cache_file = os.path.join(
            self.DUMP_FOLDERS[side + ('_gui' if is_gui else '')],
            f'{class_name}.d.ts'
        )

        imports: List[str] = deepcopy(self.imports[side])
        if class_name in imports:
            imports.remove(class_name)

        class_started = self.get_class_begin(class_name,
                                             '\n'.join(class_templates))

        text = (
                self.FILE_STARTER +
                FilterFunctionSave.generate_imports(imports, (
                    '../' if is_gui else '') + '../structure') +
                class_started
        )
        if fields:
            text += self.shift_lines('\n\n'.join(fields), 4) + '\n\n'
        if methods:
            text += self.shift_lines('\n\n'.join(methods), 4) + '\n'
        text += self.get_class_end()

        with open(cache_file, 'w', encoding='UTF-8', newline='\n') as cache:
            cache.write(text)

    def apply(self):
        for key in self.DUMP_FOLDERS:
            folder = self.DUMP_FOLDERS[key]
            if not os.path.exists(folder):
                os.mkdir(folder)

        # TODO: move into Context method
        keys = set(self.context.declarations.oop_methods.keys())
        keys.update(set(self.context.declarations.oop_fields.keys()))

        for key in sorted(keys):
            for side in ['client', 'server']:
                self.save_file_category(
                    class_name=key,
                    side=side,
                    fields=self.context.declarations.oop_fields[key].get(
                        side,
                        []
                    ),
                    methods=self.context.declarations.oop_methods[key].get(
                        side,
                        []
                    ),
                    class_templates=self.context.declarations
                                        .oop_class_templates[key].get(
                                            side,
                                            []
                                        ),
                )

        print('Generated .d.ts files with OOP\u001b[0m')


import os
import pickle
import torch
import numpy

NUM_SNIPPETS = 50

class DatasetConverter:
    @staticmethod
    def convert_pickles(path: str):
        dir_list = os.listdir(path)
        for dir_file in dir_list:

            if '.pkl' in dir_file:
                print("Opening {0}".format(dir_file))
                pkl = open(path + '/' + dir_file, 'rb')
                header = pickle.load(pkl)
                dataset = pickle.load(pkl)
                pkl.close()

                properties = dir_file.split('_')
                task = properties[1]
                if task is "T0":
                    size = properties[3]
                    type = properties[4].split('.')[0]
                else:
                    size = properties[3].split('.')[0]
                    type = "train"

                print("Converting...")
                torch_data = []
                torch_snippet = []

                if not isinstance(dataset[0], torch.Tensor):
                    # convert to torch tensors
                    images = torch.from_numpy(numpy.array(dataset[0])).permute(0, 2, 3, 1)
                    labels = torch.from_numpy(numpy.array(dataset[1]))
                else:
                    images = dataset[0]
                    labels = dataset[1]

                torch_data.append(images)
                torch_data.append(labels)
                part_images = images[:NUM_SNIPPETS].clone()
                part_labels = labels[:NUM_SNIPPETS].clone()
                torch_snippet.append(part_images)
                torch_snippet.append(part_labels)

                torch.save(torch_data, path + '/' + "SE_" + task + "_" + size + "_" + type + "_" + "full.set")
                torch.save(torch_snippet, path + '/' + "SE_" + task + "_" + size + "_" + type + "_" + "snippet.set")

                os.remove(path + '/' + dir_file)

    @staticmethod
    def change_snippet_size(path: str):
        dir_list = os.listdir(path)
        for dir_file in dir_list:
            if '.set' in dir_file and 'snippet' in dir_file:
                print("Removed.")
                os.remove(path + dir_file)

        for dir_file in dir_list:
            if '.set' in dir_file:
                print("Opening {0}".format(dir_file))
                dataset = torch.load(path + dir_file)

                properties = dir_file.split('_')
                task = properties[1]
                size = properties[2]
                type = properties[3]
                part = properties[4].split('.')[0]

                print("Converting...")
                images = dataset[0]
                labels = dataset[1]
                instance_ids = dataset[2]
                examples = dataset[3]

                # create a snippet containing NUM_SNIPPETS trajectories:
                last_id = torch.sum(examples[:NUM_SNIPPETS])
                torch_snippet = [images[:last_id].clone(),
                                 labels[:last_id].clone(),
                                 instance_ids[:last_id].clone(),
                                 examples[:NUM_SNIPPETS].clone()]

                torch.save(torch_snippet, path + "SE_" + task + "_" + size + "_" + type + "_" + "snippet.set")

if __name__ == '__main__':
    DatasetConverter.change_snippet_size('./data/eval/')
    print("Done.")


from models.user.decorator import require_admin
from flask import Blueprint, render_template, request, redirect, url_for
from werkzeug.utils import redirect
from models.store import Store
from models.user import require_admin
import json

store_blueprint = Blueprint("stores", __name__)

@store_blueprint.route("/")
def index():
    stores = Store.all()
    return render_template("stores/index.html", stores = stores)

@store_blueprint.route("/new", methods=['GET', 'POST'])
@require_admin
def new_store():
    if request.method == 'POST':
        url_prefix = request.form['url_pre']
        name = request.form['name']
        tag = request.form['tag']
        query = json.loads(request.form['query'])
        img = request.form['img']
        description = request.form['description']

        Store(url_prefix, name, tag, query, img, description).save_to_mongo()
        return redirect(url_for(".index"))

    return render_template("stores/new_store.html")

@store_blueprint.route("/edit/<string:store_id>", methods=['GET', 'POST'])
@require_admin
def edit_store(store_id):
    store = Store.get_by_id(store_id)
    if request.method == 'POST':
        store.url_prefix = request.form['url_pre']
        store.tag = request.form['tag']
        store.query = json.loads(request.form['query'])
        store.img = request.form['img']
        store.description = request.form['description']

        store.save_to_mongo()
        return redirect(url_for(".index"))
    return render_template("stores/edit_store.html", store=store)

@store_blueprint.route("/remove/<string:store_id>")
@require_admin
def remove_store(store_id):
    store = Store.get_by_id(store_id)
    store.remove_from_mongo()
    return redirect(url_for(".index"))

@store_blueprint.route("/show/<string:store_id>")
def show(store_id):
    store = Store.get_by_id(store_id)
    return render_template("stores/show.html", store=store)


import socket
import multiprocessing as mp

def handle_connection(connection, address, max_packet_size, forward_address, forward_port):
    while True:
        data = connection.recv(max_packet_size)
        if data:
            print('Data received:', data)
            new_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            new_sock.connect((forward_address, forward_port))

            new_sock.send(data)
            response = new_sock.recv(max_packet_size)
            connection.sendall(response)
        else:
            print('No more data from', address)
            connection.close()
            break

def create_proxy_server(forward_address, forward_port, max_packet_size=32768):
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
    sock.bind(('localhost', 8000))
    sock.listen(1)

    try:
        while True:
            connection, address = sock.accept()
            print('Connection from', address)
            process = mp.Process(target=handle_connection, args=(connection, address, max_packet_size, forward_address, forward_port))
            process.start()

    except KeyboardInterrupt:
        print('Shutting down...')
        sock.close()

if __name__ == '__main__':
    create_proxy_server('www.google.com', 80)


import pandas as pd
import yfinance as yf
import matplotlib.pyplot as plt

#Bringing in sentiment data
df = pd.read_csv('pfizer_dated_sentiments_prices.csv')
print(df.head())

#Bringing in stock data
historical_pfe = yf.download('PFE')
print(historical_pfe.head())

#Aligning indices before merging
historical_pfe.reset_index(inplace=True)
historical_pfe = historical_pfe.rename(columns = {'Date':'date'})
historical_pfe['date']=historical_pfe['date'].astype(str)

df['date']=df['date'].astype(str)
df.set_index('date')

#Merging
merged = pd.merge(df, historical_pfe)
print(merged)

#Examining correlations
print(merged.corr(method='pearson'))

#Rescaling prior to plotting
from sklearn.preprocessing import StandardScaler
sc = StandardScaler()

merged_ints = merged.iloc[:, 1:]
ints_scaled = pd.DataFrame(sc.fit_transform(merged_ints))
print(ints_scaled.corr(method='pearson'))

#Plotting data
date = merged['date']
closing_price = ints_scaled[6]
pos_score = ints_scaled[1]
neg_score = ints_scaled[2]
compound = ints_scaled[0]

plt.plot(date, closing_price, label="closing price", linestyle="-")
plt.plot(date, neg_score, label="negative sentiments", color='red', linestyle='--')
plt.plot(date, pos_score, label="positive sentiments", color='cyan', linestyle='--')
plt.plot(date, compound, label="compound score", color='green', linestyle="--")
plt.legend()
plt.show()


#!/usr/local/bin/python3
# -*- coding: utf-8 -*-
#
#main class
########Make snp sequences
#
#Description: program receives as input snp name and size, process the information
#and return as output sequences
# with snps in middle position
#com as combinações das snps se necessário. Está modificado para python3
#snps assembly:
#
########Algoritmo:
#1)Definir as variaveis
#2)Pegar informações da snp atraves do servidor e colocar a snp na lista
#3)Organizar a lista com base no local
#4)Fazer um loop na lista para ver se existe uma distancia maenor de 50 nts entre uma snp e outra
#5)Se não tiver, fazer as combinacões com uma snp e se tiver juntas snps e fazer as sequencias com as snps e as suas
#combinações
#
########O que ainda falta ser feito:
#1) Fazer tratamento de erros diversos
#
#
#
from snp import Snp
from allele import Allele
from graphDAO import GraphDAO
from snpDAO import SnpDAO
from sequencia import Sequence
from epigenome import Epigenome
import pandas as pd
import numpy as np
import requests, sys
import argparse

#----------------------------------------------------Métodos-----------------------------------------------------///

#----------------------------------------------Programa começa aqui----------------------------------------------///

def main(args):

	#cria uma lista vazia de snps
	#lista_snp = []

	#declara o servidor
	#server = "http://rest.ensembl.org"

	#usuario define o tamanho da lista
	#lista_tamanho = int(input("Defina o tamanho da sua lista(25 ou 50): "))

	#Usuário insere quantas snps ele que fazer as sequencia
	#numero_de_snps = int(input("coloque com quantas snps você deseja fazer as sequencias: "))

	#lista de minor allele_list
	#minor_allele_list = []

	#----------------------Isso aqui dentro de um while para poder pegar as várias snps do usuario------------------///

	#for i in range(numero_de_snps):
	genome_version = args['genome_version']

	lista_snp = args['snp_list']

	#allele1 = Allele(nome='C',local=6704583,cromossomo=1,is_comum=True,snp_pos=0)
	#allele2 = Allele(nome='A',local=6704583,cromossomo=1,is_comum=False,snp_pos=0)
	#allele7 = Allele(nome='T',local=6704583,cromossomo=1,is_comum=False,snp_pos=0)
	#allele3 = Allele(nome='T',local=6704603,cromossomo=1,is_comum=True,snp_pos=0)
	#allele4 = Allele(nome='C',local=6704603,cromossomo=1,is_comum=False,snp_pos=0)
	#allele5 = Allele(nome='T',local=135536917,cromossomo=5,is_comum=True,snp_pos=0)
	#allele6 = Allele(nome='C',local=135536917,cromossomo=5,is_comum=False,snp_pos=0)
	#allele8 = Allele(nome='G',local=85215584,cromossomo=1,is_comum=True,snp_pos=0)
	#allele9 = Allele(nome='A',local=85215584,cromossomo=1,is_comum=False,snp_pos=0)
	#allele10 = Allele(nome='C',local=85215584,cromossomo=1,is_comum=False,snp_pos=0)

	#snp1 = Snp(name='rs278020',location=6704583,chrom=1,charact='snv',ancestral_al=allele1,
				#minor_al=[allele2,allele7])
	#snp2 = Snp(name='rs278021',location=6704603,chrom=1,charact='snv',ancestral_al=allele3,
				#minor_al=[allele4])
	#snp3 = Snp(name='rs123121',location=135536917,chrom=5,charact='snv',ancestral_al=allele5,
				#minor_al=[allele6])
	#snp4 = Snp(name='rs12117219',location=85215584,chrom=1,charact='snv',ancestral_al=allele8,
				#minor_al=[allele9,allele10])

	#lista_snp = [snp1,snp2,snp3,snp4]

	#---------------------sort list by local and chromossome-------------------------///

	# timsort on chromossome and location
	lista_snp.sort(key=lambda x: (x.chrom, x.location), reverse=False)

	#-----------------------------------------------------------------verificar o tamanho----------------------------------------------------///

	lista_comb = []

	snp_stuff = SnpDAO()

	first_write = True
	#create snps dataframes
	columns_df = ['Name',
				'Chromossome',
				'Location',
				'Characteristic',
				'Allele Wild Type',
				'Allele Variation']

	#insert snp information into dataframe
	snp_df = pd.DataFrame(([i.name,
							i.chrom,
							i.location,
							i.charact,
							i.ancestral_al,
							i.minor_al] for i in lista_snp), columns=columns_df)

	#make sequences
	#group snps by their chromossomes
	snp_df_2 = snp_df.groupby('Chromossome')
	#divide them in chromossome groups
	groups = dict(list(snp_df_2))
	#iterate chromossomes groups
	for chrom_id in groups:
		#get dataframe for whole group
		chrom_df = snp_df_2.get_group(chrom_id)
		#calculate difference between locations
		chrom_df['Difference'] = chrom_df['Location'].diff()
		#null values is zero
		chrom_df.fillna(0,inplace = True)
		#get size until the end of iterations
		size = len(chrom_df)-1
		#make sequence if there is only one element
		if (size == 0):
			snp = Snp(	name=chrom_df.iloc[0]['Name'],
						location=chrom_df.iloc[0]['Location'],
						chrom=chrom_df.iloc[0]['Chromossome'],
						charact=chrom_df.iloc[0]['Characteristic'],
						ancestral_al=chrom_df.iloc[0]['Allele Wild Type'],
						minor_al=chrom_df.iloc[0]['Allele Variation'])
			first_write = snp_stuff.request_sequence(snp,genome_version,first_write)
		#check if there is more than one element
		elif (size > 0):
			#iterate chromossome list
			for i in range(len(chrom_df)-1):
				next_smaller = chrom_df.iloc[i + 1]['Difference'] <31
				bigger = chrom_df.iloc[i]['Difference'] >=31
				current_0 = chrom_df.iloc[i]['Difference'] == 0
				smaller = chrom_df.iloc[i]['Difference'] <31
				#check if it needs to be in list of combinations
				if ( (bigger and next_smaller) or (smaller and not(current_0)) or (current_0 and next_smaller) ):
					#TODO see if it works. Added now
					if (len(lista_comb) > 1 ):
						if (bigger and next_smaller):
							first_write = snp_stuff.request_sequence_combinations(lista_comb,genome_version,first_write)
                    		#delete combinations list
							del lista_comb[:]

					snp = Snp(	name=chrom_df.iloc[i]['Name'],
								location=chrom_df.iloc[i]['Location'],
								chrom=chrom_df.iloc[i]['Chromossome'],
								charact=chrom_df.iloc[i]['Characteristic'],
								ancestral_al=chrom_df.iloc[i]['Allele Wild Type'],
								minor_al=chrom_df.iloc[i]['Allele Variation'])
					lista_comb.append(snp)
					#if its last iteration
					if i == size-1:
						if (next_smaller):
							snp = Snp(	name=chrom_df.iloc[i+1]['Name'],
										location=chrom_df.iloc[i+1]['Location'],
										chrom=chrom_df.iloc[i+1]['Chromossome'],
										charact=chrom_df.iloc[i+1]['Characteristic'],
										ancestral_al=chrom_df.iloc[i+1]['Allele Wild Type'],
										minor_al=chrom_df.iloc[i+1]['Allele Variation'])
							lista_comb.append(snp)
							first_write = snp_stuff.request_sequence_combinations(lista_comb,genome_version,first_write)
							#delete combinations list
							del lista_comb[:]
						#last iteration is not in combinations list
						else:
							snp = Snp(	name=chrom_df.iloc[i+1]['Name'],
										location=chrom_df.iloc[i+1]['Location'],
										chrom=chrom_df.iloc[i+1]['Chromossome'],
										charact=chrom_df.iloc[i+1]['Characteristic'],
										ancestral_al=chrom_df.iloc[i+1]['Allele Wild Type'],
										minor_al=chrom_df.iloc[i+1]['Allele Variation'])
							print("176 name: " + chrom_df.iloc[i+1]['Name'])
							first_write = snp_stuff.request_sequence(snp,genome_version,first_write)
							if (len(lista_comb) > 1 ):
								first_write = snp_stuff.request_sequence_combinations(lista_comb,genome_version,first_write)
                    			#delete combinations list
								del lista_comb[:]
				#if its not in combinations list make sequence
				else:
					#if there is more than one in combinations list make sequence
					if (len(lista_comb) > 1 ):
						print("------DEBUG------")
						first_write = snp_stuff.request_sequence_combinations(lista_comb,genome_version,first_write)
                    	#delete combinations list
						del lista_comb[:]
					snp = Snp(	name=chrom_df.iloc[i]['Name'],
								location= chrom_df.iloc[i]['Location'],
								chrom=chrom_df.iloc[i]['Chromossome'],
								charact=chrom_df.iloc[i]['Characteristic'],
								ancestral_al=chrom_df.iloc[i]['Allele Wild Type'],
								minor_al=chrom_df.iloc[i]['Allele Variation'])
					print("196 name: " + chrom_df.iloc[i]['Name'])
					first_write = snp_stuff.request_sequence(snp,genome_version,first_write)
					if i == size-1:
						snp = Snp(	name=chrom_df.iloc[i+1]['Name'],
									location= chrom_df.iloc[i+1]['Location'],
									chrom=chrom_df.iloc[i+1]['Chromossome'],
									charact=chrom_df.iloc[i+1]['Characteristic'],
									ancestral_al=chrom_df.iloc[i+1]['Allele Wild Type'],
									minor_al=chrom_df.iloc[i+1]['Allele Variation'])
						print("205 name: " + chrom_df.iloc[i+1]['Name'])
						first_write = snp_stuff.request_sequence(snp,genome_version,first_write)
				# next_smaller = chrom_df.iloc[i + 1]['Difference'] <32
				# next_bigger = chrom_df.iloc[i + 1]['Difference'] >32
				# bigger = chrom_df.iloc[i]['Difference'] >=32
				# current_0 = chrom_df.iloc[i]['Difference'] == 0
				# smaller = chrom_df.iloc[i]['Difference'] <32
				# ##new logic
				# #check if it needs to be in list of combinations/ always checking the next_smaller ones
				# if (next_smaller):
				# 	snp = Snp(	name=chrom_df.iloc[i]['Name'],
				# 				location=chrom_df.iloc[i]['Location'],
				# 				chrom=chrom_df.iloc[i]['Chromossome'],
				# 				charact=chrom_df.iloc[i]['Characteristic'],
				# 				ancestral_al=chrom_df.iloc[i]['Allele Wild Type'],
				# 				minor_al=chrom_df.iloc[i]['Allele Variation'])
				# 	lista_comb.append(snp)
				# 	#if its last iteration
				# 	if (i == size-1):
				# 		snp = Snp(	name=chrom_df.iloc[i + 1]['Name'],
				# 					location=chrom_df.iloc[i + 1]['Location'],
				# 					chrom=chrom_df.iloc[i + 1]['Chromossome'],
				# 					charact=chrom_df.iloc[i + 1]['Characteristic'],
				# 					ancestral_al=chrom_df.iloc[i + 1]['Allele Wild Type'],
				# 					minor_al=chrom_df.iloc[i + 1]['Allele Variation'])
				# 		lista_comb.append(snp)
				# 		first_write = snp_stuff.request_sequence_combinations(lista_comb,genome_version,first_write)
				# 		del lista_comb[:]
				# #if next one is not smaller
				# else:
				# 	if (smaller and not(current_0)):
				# 		snp = Snp(	name=chrom_df.iloc[i]['Name'],
				# 				location=chrom_df.iloc[i]['Location'],
				# 				chrom=chrom_df.iloc[i]['Chromossome'],
				# 				charact=chrom_df.iloc[i]['Characteristic'],
				# 				ancestral_al=chrom_df.iloc[i]['Allele Wild Type'],
				# 				minor_al=chrom_df.iloc[i]['Allele Variation'])
				# 		lista_comb.append(snp)
				# 		first_write = snp_stuff.request_sequence_combinations(lista_comb,genome_version,first_write)
				# 		del lista_comb[:]
				# 		if (i == size-1):
				# 			snp = Snp(	name=chrom_df.iloc[i + 1]['Name'],
				# 				location=chrom_df.iloc[i + 1]['Location'],
				# 				chrom=chrom_df.iloc[i + 1]['Chromossome'],
				# 				charact=chrom_df.iloc[i + 1]['Characteristic'],
				# 				ancestral_al=chrom_df.iloc[i + 1]['Allele Wild Type'],
				# 				minor_al=chrom_df.iloc[i + 1]['Allele Variation'])
				# 			first_write = snp_stuff.request_sequence(snp,genome_version,first_write)
				# 	else:
				# 		snp = Snp(	name=chrom_df.iloc[i]['Name'],
				# 				location=chrom_df.iloc[i]['Location'],
				# 				chrom=chrom_df.iloc[i]['Chromossome'],
				# 				charact=chrom_df.iloc[i]['Characteristic'],
				# 				ancestral_al=chrom_df.iloc[i]['Allele Wild Type'],
				# 				minor_al=chrom_df.iloc[i]['Allele Variation'])
				# 		first_write = snp_stuff.request_sequence(snp,genome_version,first_write)

				##old logic
				#check if it needs to be in list of combinations
				# if ( (bigger and next_smaller) or (smaller and not(current_0)) or (current_0 and next_smaller) ):

				# 	snp = Snp(	name=chrom_df.iloc[i]['Name'],
				# 				location=chrom_df.iloc[i]['Location'],
				# 				chrom=chrom_df.iloc[i]['Chromossome'],
				# 				charact=chrom_df.iloc[i]['Characteristic'],
				# 				ancestral_al=chrom_df.iloc[i]['Allele Wild Type'],
				# 				minor_al=chrom_df.iloc[i]['Allele Variation'])
				# 	lista_comb.append(snp)
				# 	#if its last iteration
				# 	if i == size-1:
				# 		if (next_smaller):
				# 			snp = Snp(	name=chrom_df.iloc[i+1]['Name'],
				# 						location=chrom_df.iloc[i+1]['Location'],
				# 						chrom=chrom_df.iloc[i+1]['Chromossome'],
				# 						charact=chrom_df.iloc[i+1]['Characteristic'],
				# 						ancestral_al=chrom_df.iloc[i+1]['Allele Wild Type'],
				# 						minor_al=chrom_df.iloc[i+1]['Allele Variation'])
				# 			lista_comb.append(snp)
				# 			first_write = snp_stuff.request_sequence_combinations(lista_comb,genome_version,first_write)
				# 			#delete combinations list
				# 			del lista_comb[:]
				# 		#last iteration is not in combinations list
				# 		else:
				# 			snp = Snp(	name=chrom_df.iloc[i+1]['Name'],
				# 						location=chrom_df.iloc[i+1]['Location'],
				# 						chrom=chrom_df.iloc[i+1]['Chromossome'],
				# 						charact=chrom_df.iloc[i+1]['Characteristic'],
				# 						ancestral_al=chrom_df.iloc[i+1]['Allele Wild Type'],
				# 						minor_al=chrom_df.iloc[i+1]['Allele Variation'])
				# 			first_write = snp_stuff.request_sequence(snp,genome_version,first_write)
				# 			if (len(lista_comb) > 1 ):
				# 				first_write = snp_stuff.request_sequence_combinations(lista_comb,genome_version,first_write)
                #     			#delete combinations list
				# 				del lista_comb[:]
				# if (smaller and next_bigger):
				# 	if (len(lista_comb) > 1 ):
				# 		print("------DEBUG1------")
				# 		first_write = snp_stuff.request_sequence_combinations(lista_comb,genome_version,first_write)
                #     	#delete combinations list
				# 		del lista_comb[:]
				# #if its not in combinations list make sequence
				# else:
				# 	#if there is more than one in combinations list make sequence
				# 	if (len(lista_comb) > 1 ):
				# 		print("------DEBUG2------")
				# 		first_write = snp_stuff.request_sequence_combinations(lista_comb,genome_version,first_write)
                #     	#delete combinations list
				# 		del lista_comb[:]
				# 	snp = Snp(	name=chrom_df.iloc[i]['Name'],
				# 				location= chrom_df.iloc[i]['Location'],
				# 				chrom=chrom_df.iloc[i]['Chromossome'],
				# 				charact=chrom_df.iloc[i]['Characteristic'],
				# 				ancestral_al=chrom_df.iloc[i]['Allele Wild Type'],
				# 				minor_al=chrom_df.iloc[i]['Allele Variation'])
				# 	first_write = snp_stuff.request_sequence(snp,genome_version,first_write)
				# 	if i == size-1:
				# 		snp = Snp(	name=chrom_df.iloc[i+1]['Name'],
				# 					location= chrom_df.iloc[i+1]['Location'],
				# 					chrom=chrom_df.iloc[i+1]['Chromossome'],
				# 					charact=chrom_df.iloc[i+1]['Characteristic'],
				# 					ancestral_al=chrom_df.iloc[i+1]['Allele Wild Type'],
				# 					minor_al=chrom_df.iloc[i+1]['Allele Variation'])
				# 	first_write = snp_stuff.request_sequence(snp,genome_version,first_write)
	lista_comb = []

	return txt_2_list(args['return_list'], filename = "sequenciasdef.fna")

def txt_2_list(return_list, filename = "sequenciasdef.fna"):
  if(return_list):
    with open(filename,"r") as f:
      txt_to_list =[]
      for line in f:
        txt_to_list.append(line)
    f.close
    return txt_to_list
  else:
    return False

if __name__ == '__main__':

	parser = argparse.ArgumentParser(description='')
	parser.add_argument('--snp_list', default=[], type=list)
	parser.add_argument('--verify', default='n', type=str)
	parser.add_argument('--delete_snp', default='n', type=str)
	parser.add_argument('--return_list', default=False, type=bool)
	parser.add_argument('--genome_version', default='GRCh37', type=str)
	args = vars(parser.parse_args())
	main(args)


import torch
from torch import nn
from torchvision.transforms import InterpolationMode
import torchvision.transforms.functional as TF
import timm
from .layers import AttentionRefinementModule, FeatureFusionModule, ASPPv2, Spade, C3Blk, ResBlk, SPP, Conv, SpatialAttentionRefinementModule, PyramidPoolAgg, SIM, SegHead
from .metaformer import *

class NoIMUDecoder(nn.Module):
    """Decoder without IMU input."""
    def __init__(self, num_classes=3):
        super(NoIMUDecoder, self).__init__()

        self.arm1 = AttentionRefinementModule(2048)
        self.arm2 = nn.Sequential(
            AttentionRefinementModule(512, last_arm=True),
            nn.Conv2d(512, 2048, 1) # Equalize number of features with ARM1
        )

        self.ffm = FeatureFusionModule(256, 2048, 1024)
        self.aspp = ASPPv2(1024, [6, 12, 18, 24], num_classes)

    def forward(self, x, aux, skip2, skip1, imu_mask):

        arm1 = self.arm1(x)
        arm2 = self.arm2(skip2)
        arm_combined = arm1 + arm2

        x = self.ffm(skip1, arm_combined)

        output = self.aspp(x)

        return output

class IMUDecoder(nn.Module):
    """Decoder with IMU information merging."""
    def __init__(self, num_classes=3):
        super(IMUDecoder, self).__init__()

        self.arm1 = AttentionRefinementModule(2048 + 1)
        self.aspp1 = ASPPv2(2048, [6, 12, 18], 32)
        self.ffm1 = FeatureFusionModule(2048 + 1, 32, 1024)

        self.arm2 = nn.Sequential(
            AttentionRefinementModule(512 + 1, last_arm=True),
            nn.Conv2d(512 + 1, 1024, 1, bias=False) # Equalize number of features with FFM1
        )

        self.ffm = FeatureFusionModule(256 + 1, 1024, 1024)
        self.aspp = ASPPv2(1024, [6, 12, 18, 24], num_classes)
        self.upsample1 = nn.Upsample(scale_factor=2.0, mode='bilinear', align_corners=True)

    def forward(self, out, aux, skip2, skip1, imu_mask):

        # resize IMU mask to two required scales
        imu_mask_s1 = TF.resize(imu_mask, (out.size(2), out.size(3)), InterpolationMode.NEAREST)
        imu_mask_s0 = TF.resize(imu_mask, (skip1.size(2), skip1.size(3)), InterpolationMode.NEAREST)

        # feature mixer
        out_imu = torch.cat([out, imu_mask_s1], dim=1)
        arm1 = self.arm1(out_imu)
        aspp1 = self.aspp1(out)
        ffm1 = self.ffm1(arm1, aspp1)

        # enriched skip connection
        skip2_imu = torch.cat([skip2, imu_mask_s1], dim=1)
        arm2 = self.arm2(skip2_imu)

        # decoder
        arm_combined = ffm1 + arm2
        skip1_imu = torch.cat([skip1, imu_mask_s0], dim=1)
        x = self.ffm(skip1_imu, arm_combined)

        # head
        output = self.aspp(x)

        return output

class IMUDecoderSmall(nn.Module):
    """WaSR-light IMU decoder."""
    def __init__(self, num_classes=3, ch = [512, 256, 128, 64]):
        super(IMUDecoderSmall, self).__init__()

        self.arm1 = AttentionRefinementModule(ch[0] + 1)
        self.aspp1 = ASPPv2(ch[0], [6, 12, 18], 32)
        self.ffm1 = FeatureFusionModule(ch[0] + 1, 32, ch[1])

        self.arm2 = nn.Sequential(
                AttentionRefinementModule(ch[2] + 1, last_arm=True),
                nn.Conv2d(ch[2] + 1, ch[1], 1, bias=False) # Equalize number of features with FFM1
            )

        self.ffm = FeatureFusionModule(ch[3] + 1, ch[1], ch[1])
        self.aspp = ASPPv2(ch[1], [6, 12, 18, 24], num_classes)

        self.upsample= nn.Upsample(scale_factor=4.0, mode='bilinear', align_corners=True)
        self.upsample1 = nn.Upsample(scale_factor=2.0, mode='bilinear', align_corners=True)

    def forward(self, out, aux, skip2, skip1, imu_mask):

        # resize IMU mask to the required scales
        imu_mask_s2 = TF.resize(imu_mask, (out.size(2), out.size(3)), InterpolationMode.NEAREST)
        imu_mask_s1 = TF.resize(imu_mask, (skip2.size(2), skip2.size(3)), InterpolationMode.NEAREST)
        imu_mask_s0 = TF.resize(imu_mask, (skip1.size(2), skip1.size(3)), InterpolationMode.NEAREST)

        # feature mixer
        out_imu = torch.cat([out, imu_mask_s2], dim=1)
        arm1 = self.arm1(out_imu)
        aspp1 = self.aspp1(out)
        ffm1 = self.ffm1(arm1, aspp1)

        # enriched skip connection
        skip2_imu = torch.cat([skip2, imu_mask_s1], dim=1)
        arm2 = self.arm2(skip2_imu)

        # decoder
        ffm1_up = self.upsample(ffm1)
        arm_combined = ffm1_up + arm2
        skip1_imu = torch.cat([skip1, imu_mask_s0], dim=1)
        arm_combined_up = self.upsample1(arm_combined)
        x = self.ffm(skip1_imu, arm_combined_up)

        # head
        output = self.aspp(x)

        return output

def get_token_mixer(letter):
    if letter == "A":
        return Attention
    elif letter == "P":
        return Pooling
    elif letter == "C":
        return ARMMixer
    elif letter == "S":
        return SpatialAttentionMixer
    else:
        raise ValueError(f"Mixer {mixer} not supported.")

class EWaSRDecoder(nn.Module):

    def __init__(self, num_classes=3, L = 6, ch=512, ch_sim=256, mixer="CCCCSS", enricher="SS", imu=False, project=False):

        super(EWaSRDecoder, self).__init__()

        self.project = project
        self.ch = [ch, ch//2, ch//4, ch//8] if isinstance(ch, int) else ch

        if self.project:
            self.convs_project1 = nn.Conv2d(self.ch[0], self.ch[0] // 2, 1)
            self.convs_project2 = nn.Conv2d(self.ch[1], self.ch[1] // 2, 1)
            self.convs_project3 = nn.Conv2d(self.ch[2], self.ch[2] // 2, 1)
            self.convs_project4 = nn.Conv2d(self.ch[3], self.ch[3] // 2, 1)
            self.ch = [ch // 2 for ch in self.ch]

        self.imu = imu
        self.pool_feats = PyramidPoolAgg(2)
        self.metaformers = None
        self.metaformers_skip2 = None

        if len(mixer) > 0:
            metaformers = []
            for letter in mixer:
                metaformers.append(MetaFormerBlock(sum(self.ch), token_mixer=get_token_mixer(letter)))
            self.metaformers = nn.Sequential(*metaformers)

        if len(enricher) > 0:
            metaformers_skip2 = []
            for letter in enricher:
                metaformers_skip2.append(MetaFormerBlock(self.ch[2], token_mixer=get_token_mixer(letter)))
            self.metaformers_skip2 = nn.Sequential(*metaformers_skip2)

        """
        if mix:
            metaformers = [(
                MetaFormerBlock(sum(self.ch), token_mixer=ARMMixer),
                MetaFormerBlock(sum(self.ch), token_mixer=SpatialAttentionMixer))
                 for _ in range(L//2)]
            metaformers = [item for sublist in metaformers for item in sublist]
        elif short is None:
            metaformers = [MetaFormerBlock(sum(self.ch), token_mixer=token_mixer) for _ in range(L)]
        else:
            metaformers = [MetaFormerBlock(sum(self.ch), token_mixer=token_mixer) for _ in range(L-2)]
            short_mixer = SpatialAttentionMixer if short == "sarm" else ARMMixer
            metaformers.extend([
                MetaFormerBlock(sum(self.ch), token_mixer=short_mixer),
                MetaFormerBlock(sum(self.ch), token_mixer=short_mixer)
            ])
        self.metaformers = nn.Sequential(*metaformers)

        if mix:
            metaformers_skip2 = [MetaFormerBlock(self.ch[2], token_mixer=ARMMixer), MetaFormerBlock(self.ch[2], token_mixer=SpatialAttentionMixer)]
            self.metaformers_skip2 = nn.Sequential(*metaformers_skip2)
        elif skip is not None:
            token_mixer_skip2 = SpatialAttentionMixer if skip == "sarm" else ARMMixer
            metaformers_skip2 = [MetaFormerBlock(self.ch[2], token_mixer=token_mixer_skip2) for _ in range(2)]
            self.metaformers_skip2 = nn.Sequential(*metaformers_skip2)
        """

        self.sim1 = SIM(self.ch[0], ch_sim)
        self.sim2 = SIM(self.ch[1], ch_sim)
        self.sim3 = SIM(self.ch[2], ch_sim)
        self.sim4 = SIM(self.ch[3], ch_sim)

        head_in = ch_sim+1 if self.imu else ch_sim
        self.seg_head = SegHead(head_in, num_classes)

    def forward(self, x, aux, skip2, skip1, imu_mask):

        if self.project:
            x = self.convs_project1(x)
            aux = self.convs_project2(aux)
            skip2 = self.convs_project3(skip2)
            skip1 = self.convs_project4(skip1)

        tokens = self.pool_feats([x, aux, skip2, skip1])

        if self.metaformers is not None:
            tokens = self.metaformers(tokens)

        if self.metaformers_skip2 is not None:
            skip2 = self.metaformers_skip2(skip2)

        f1 = self.sim1(x, tokens[:, :self.ch[0], :, :])
        f2 = self.sim2(aux, tokens[:, self.ch[0]:sum(self.ch[:2]), :, :])
        f3 = self.sim3(skip2, tokens[:, sum(self.ch[:2]):sum(self.ch[:3]), :, :])
        f4 = self.sim4(skip1, tokens[:, sum(self.ch[:3]):, :, :])

        f1 = TF.resize(f1, (f4.size(2), f4.size(3)), InterpolationMode.BILINEAR)
        f2 = TF.resize(f2, (f4.size(2), f4.size(3)), InterpolationMode.BILINEAR)
        f3 = TF.resize(f3, (f4.size(2), f4.size(3)), InterpolationMode.BILINEAR)

        x = f1 + f2 + f3 + f4

        if self.imu:
            imu_mask = TF.resize(imu_mask, (x.size(2), x.size(3)), InterpolationMode.NEAREST)
            x = torch.cat([x, imu_mask], dim=1)

        x = self.seg_head(x)

        return x


from sentence_transformers import SentenceTransformer
import scipy.spatial
from datetime import datetime

def cal_sim(wn_des_list, wiki_des_list, model_path, batch_size):
    start = datetime.now()
    print('start time:', start.strftime("%Y-%m-%d %H:%M:%S"), '\n==================')
    print('model path:', model_path)
    print('synset number:', len(wn_des_list))
    embedder = SentenceTransformer(model_path)
    wiki_len = [len(i) for i in wiki_des_list]
    wiki_des_list = sum(wiki_des_list, [])
    # handling 'None' to prevent error
    wiki_des_list = [i if i != 'None' else 'None Description' for i in wiki_des_list]

    wiki_embeddings = embedder.encode(wiki_des_list, batch_size=batch_size)
    wn_embeddings = embedder.encode(wn_des_list, batch_size=batch_size)
    print(len(wiki_embeddings), len(wn_embeddings))

    wiki_embeddings_lst = []
    for index, i in enumerate(wiki_len):
        wiki_embeddings_lst.append(wiki_embeddings[sum(wiki_len[:index]):sum(wiki_len[:index+1])])
    assert len(wiki_embeddings_lst) == len(wn_embeddings)
    print(len(wiki_embeddings_lst), len(wn_embeddings))

    all_sim_lst = []
    for wn_emb, wiki_emb in zip(wn_embeddings, wiki_embeddings_lst):
        sim_lst = []
        distances = scipy.spatial.distance.cdist([wn_emb], wiki_emb, "cosine")[0]
        results = zip(range(len(distances)), distances)

        for idx, distance in results:
            sim = round(1 - distance, 6)
            sim_lst.append(sim)
        all_sim_lst.append(sim_lst)

    end = datetime.now()
    print('==================\nend time:', end.strftime("%Y-%m-%d %H:%M:%S"))
    print('running time is', round((end-start).seconds/60, 4), 'minutes')
    return all_sim_lst


"""SPADE Localizer Class."""

import logging
from typing import Optional, Tuple, Union

import numpy as np
from pyscf import dft, gto
from scipy import linalg

from .base import Localizer

logger = logging.getLogger(__name__)

class SPADELocalizer(Localizer):
    """Object used to localise molecular orbitals (MOs) using SPADE Localization.

    Running localization returns active and environment systems.

    Args:
        pyscf_rks (gto.Mole): PySCF molecule object
        n_active_atoms (int): Number of active atoms
        localization_method (str): String of orbital localization method (spade, pipekmezey, boys, ibo)
        occ_cutoff (float): Threshold for selecting occupied active region (only requried if
                                spade localization is NOT used)
        virt_cutoff (float): Threshold for selecting unoccupied (virtual) active region (required for
                                spade approach too!)
        run_virtual_localization (bool): optional flag on whether to perform localization of virtual orbitals.
                                         Note if False appends canonical virtual orbs to C_loc_occ_and_virt matrix

    Attributes:
        c_active (np.array): C matrix of localized occupied active MOs (columns define MOs)
        c_enviro (np.array): C matrix of localized occupied ennironment MOs
        c_loc_occ_and_virt (np.array): Full localized C_matrix (occpuied and virtual)
        dm_active (np.array): active system density matrix
        dm_enviro (np.array): environment system density matrix
        active_MO_inds (np.array): 1D array of active occupied MO indices
        enviro_MO_inds (np.array): 1D array of environment occupied MO indices
        _c_loc_occ (np.array): C matrix of localized occupied MOs

    Methods:
        run: Main function to run localization.
    """

    def __init__(
        self,
        pyscf_scf: gto.Mole,
        n_active_atoms: int,
        occ_cutoff: Optional[float] = 0.95,
        virt_cutoff: Optional[float] = 0.95,
        run_virtual_localization: Optional[bool] = False,
    ):
        """Initialize SPADE Localizer object."""
        super().__init__(
            pyscf_scf,
            n_active_atoms,
            occ_cutoff=occ_cutoff,
            virt_cutoff=virt_cutoff,
            run_virtual_localization=run_virtual_localization,
        )

    def _localize_spin(
        self, c_matrix: np.ndarray, occupancy: np.ndarray
    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
        """Localize orbitals of one spin using SPADE.

        Args:
            c_matrix (np.ndarray): Unlocalized C matrix of occupied orbitals.
            occupancy (np.ndarray): Occupancy of orbitals.

        Returns:
            np.ndarray: Localized C matrix of occupied orbitals.
        """
        logger.debug("Localising with SPADE.")

        # We want the same partition for each spin.
        # It wouldn't make sense to have different spin states be localized differently.

        n_occupied_orbitals = np.count_nonzero(occupancy)
        occupied_orbitals = c_matrix[:, :n_occupied_orbitals]

        n_act_aos = self._global_ks.mol.aoslice_by_atom()[self._n_active_atoms - 1][-1]
        logger.debug(f"{n_act_aos} active AOs.")

        ao_overlap = self._global_ks.get_ovlp()

        # Orbital rotation and partition into subsystems A and B
        # rotation_matrix, sigma = embed.orbital_rotation(occupied_orbitals,
        #    n_act_aos, ao_overlap)

        rotated_orbitals = (
            linalg.fractional_matrix_power(ao_overlap, 0.5) @ occupied_orbitals
        )
        _, sigma, right_vectors = linalg.svd(rotated_orbitals[:n_act_aos, :])

        logger.debug(f"Singular Values: {sigma}")

        # n_act_mos, n_env_mos = embed.orbital_partition(sigma)
        # Prevents an error with argmax
        if len(sigma) == 1:
            n_act_mos = 1
        else:
            value_diffs = sigma[:-1] - sigma[1:]
            n_act_mos = np.argmax(value_diffs) + 1
        n_env_mos = n_occupied_orbitals - n_act_mos
        logger.debug(f"{n_act_mos} active MOs.")
        logger.debug(f"{n_env_mos} environment MOs.")

        # get active and enviro indices
        active_MO_inds = np.arange(n_act_mos)
        enviro_MO_inds = np.arange(n_act_mos, n_act_mos + n_env_mos)

        # Defining active and environment orbitals and density
        c_active = occupied_orbitals @ right_vectors.T[:, :n_act_mos]
        c_enviro = occupied_orbitals @ right_vectors.T[:, n_act_mos:]
        c_loc_occ = occupied_orbitals @ right_vectors.T

        # storing condition used to select env system
        self.enviro_selection_condition = sigma

        return (active_MO_inds, enviro_MO_inds, c_active, c_enviro, c_loc_occ)


import numpy as np

from utils.params import AgentActionType, UserActionType

def deco(cls):
    cls._define_function()
    return cls

@deco
class UserFeatures(object):
    dimensions = 0
    _function = {}

    @classmethod
    def get_vector(cls, state, action):
        """Returns the feature vector for a state-action pair.

        Args:
            state (AgentActionType): User state.
            action (UserActionType): User action.

        Returns:
            numpy.array: Feature vector for the given state-action pair.
        """
        try:
            return cls._function[(state, action)]
        except KeyError as e:
            print("Exception: {}. Invalid state: '{}' and action: '{}' pair"
                  .format(e, state, action))
            raise

    @classmethod
    def _define_function(cls):
        """Builds the feature function.
        """
        # First calculate the size of input space.
        cls.dimensions = 0
        # User-state is effectively characterized by AgentActionType
        for state in AgentActionType:
            for action in UserActionType:
                cls.dimensions += 1

        # Now construct the feature function -- mapping from state-action pairs
        # to an n dim vector whose elements are in the range [0, 1]
        i = 0
        for state in AgentActionType:
            for action in UserActionType:
                vec = np.zeros(cls.dimensions)
                vec[i] = 1.
                cls._function[(state, action)] = vec
                i += 1


import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.distributions as D
import pytorch_ssim
import torchvision
from GQN.core import LSTMCore, GRUCore
from util.datasets import get_buffer_length
import copy

"""
Implementation of the generator used in the GQN paper
can also replace LSTM cells with GRU cells
"""
class GQNGenerator(nn.Module):
    def __init__(self, settings, index):
        super(GQNGenerator, self).__init__()
        self.core_count = settings.model.generators[index].settings.core_count
        self.share = settings.model.generators[index].settings.weight_sharing
        self.upscale_share = settings.model.generators[index].settings.upscale_sharing
        self.index = index

        # Set dimensions
        self.h_dim = settings.model.generators[index].settings.state_dim
        self.v_dim = settings.model.pose_dim
        self.z_dim = settings.model.generators[index].settings.latent_dim
        self.height = settings.model.generators[index].render_size
        self.width = settings.model.generators[index].render_size
        self.representations = settings.model.generators[index].representations

        self.x_dim = 0
        for buffer in settings.model.generators[index].query_passes:
            self.x_dim += get_buffer_length(buffer)

        self.r_dim = 0
        for r in self.representations:
            self.r_dim += settings.model.representations[r].representation_dim

        self.scaling = settings.model.generators[index].settings.downscaling
        self.downscaled_height = self.height // self.scaling
        self.downscaled_width = self.width // self.scaling
        self.out_dim = 0
        for buffer in settings.model.generators[index].output_passes:
            self.out_dim += get_buffer_length(buffer)

        inference_dim = self.v_dim + self.r_dim + self.out_dim + self.h_dim
        generator_dim = self.v_dim + self.r_dim + self.z_dim + self.x_dim

        if settings.model.generators[index].settings.cell_type == "GRU":
            self.core = GRUCore
            self.gru = True
        else:
            self.core = LSTMCore
            self.gru = False

        # Set up computational cores
        if self.share:
            self.inference = self.core(inference_dim, self.h_dim)
            self.generator = self.core(generator_dim, self.h_dim)
        else:
            self.inference = nn.ModuleList([self.core(inference_dim, self.h_dim) for _ in range(self.core_count)])
            self.generator = nn.ModuleList([self.core(generator_dim, self.h_dim) for _ in range(self.core_count)])

        # These convolutions are used to model the normal distributions
        self.prior = nn.Conv2d(self.h_dim, self.z_dim * 2, 5, padding=2)
        self.posterior = nn.Conv2d(self.h_dim, self.z_dim * 2, 5, padding=2)

        # Final convolution on the canvas vairable
        self.observation = nn.Conv2d(self.h_dim, self.out_dim, 1)

        # Scaling convolutions performed on input and canvas of cores
        # No bias because we don't want to add information that is not already in the input
        if self.x_dim > 0:
            self.downsample = nn.Conv2d(self.x_dim, self.x_dim, self.scaling, stride=self.scaling, padding=0, bias=False)
        if self.upscale_share:
            self.upsample = nn.ConvTranspose2d(self.h_dim, self.h_dim, self.scaling, stride=self.scaling, padding=0, bias=False)
        else:
            self.upsample = nn.ModuleList([nn.ConvTranspose2d(self.h_dim, self.h_dim, self.scaling, stride=self.scaling, padding=0, bias=False) for _ in range(self.core_count)])

    def forward(self, x, y, o, v, p, r, iteration, additional=None):
        """
        h_0, c_0, u_0 = (0, 0, 0)
        for l in L:
            compute prior factor from normal distribution
            update inference state
            compute posterior factor from normal distribution
            sample posterior to get z_l
            update generator state
            add to ELBO KL computation
        """

        kl = 0
        batch_size = y.shape[0]
        inf_state_h = torch.zeros([batch_size, self.h_dim, self.downscaled_height, self.downscaled_width], dtype=x.dtype, device=x.device)
        gen_state_h = torch.zeros([batch_size, self.h_dim, self.downscaled_height, self.downscaled_width], dtype=x.dtype, device=x.device)

        if not self.gru:
            inf_state_c = torch.zeros([batch_size, self.h_dim, self.downscaled_height, self.downscaled_width], dtype=x.dtype, device=x.device)
            gen_state_c = torch.zeros([batch_size, self.h_dim, self.downscaled_height, self.downscaled_width], dtype=x.dtype, device=x.device)

        canvas = torch.zeros([batch_size, self.h_dim, self.height, self.width], dtype=x.dtype, device=x.device)

        r = r.repeat(1, 1, self.downscaled_height // r.shape[2], self.downscaled_width // r.shape[3])
        v = v.repeat(1, 1, self.downscaled_height // v.shape[2], self.downscaled_width // v.shape[3])

        if self.x_dim > 0: x_ds = self.downsample(x)
        else: x_ds = torch.zeros([batch_size, 0, self.downscaled_height, self.downscaled_width], dtype=x.dtype, device=x.device)
        y_ds = F.interpolate(y, size=(self.downscaled_height, self.downscaled_width))

        for l in range(self.core_count):
            # Compute prior distribution by extracting mean and logstd
            p_mean, p_logstd = torch.chunk(self.prior(gen_state_h), 2, dim=1)
            prior_distr = D.Normal(p_mean, F.softplus(p_logstd))

            # Inference state update
            inf_input = torch.cat([gen_state_h, y_ds, v, r], dim=1)

            # Select cell
            if self.share:
                cell = self.inference
            else:
                cell = self.inference[l]

            if self.gru:
                inf_state_h = cell(inf_input, inf_state_h)
            else:
                inf_state_h, inf_state_c = cell(inf_input, inf_state_h, inf_state_c)
            # Compute posterior distribution by extracting mean and logstd
            q_mean, q_logstd = torch.chunk(self.posterior(inf_state_h), 2, dim=1)
            posterior_distr = D.Normal(q_mean, F.softplus(q_logstd))

            # Generator state update
            z = posterior_distr.rsample() # rsample makes sample differentiable
            gen_input = torch.cat([v, r, z, x_ds], dim=1)

            if self.share:
                cell = self.generator
            else:
                cell = self.generator[l]

            if self.gru:
                gen_state_h = cell(gen_input, gen_state_h)
            else:
                gen_state_h, gen_state_c = cell(gen_input, gen_state_h, gen_state_c)

            # Update canvas and KL divergence
            if self.upscale_share:
                canvas += self.upsample(gen_state_h)
            else:
                canvas += self.upsample[l](gen_state_h)
            kl += D.kl.kl_divergence(posterior_distr, prior_distr)

        y_mean = self.observation(canvas)

        # Compute ELBO
        sigma = max(0.7 + (2.0 - 0.7) * (1 - iteration / 2e5), 0.7)
        ll = D.Normal(y_mean, sigma).log_prob(torch.log1p(y))
        ll = torch.sum(ll, dim=[1,2,3])
        kl = torch.sum(kl, dim=[1,2,3])

        losses = {}
        losses["generator" + str(self.index) + "_kl"] = kl / (self.width * 25)
        losses["generator" + str(self.index) + "_ll"] = -ll / (self.width * 25)

        outputs = {}
        outputs["generated"] = torch.exp(y_mean) - 1
        outputs["losses"] = losses

        return outputs

    def sample(self, x, o, v, p, r):
        batch_size = r.shape[0]
        w = x.shape[2]
        dh = w // self.scaling
        gen_state_h = torch.zeros([batch_size, self.h_dim, dh, dh], dtype=v.dtype, device=v.device)

        if not self.gru:
            gen_state_c = torch.zeros([batch_size, self.h_dim, dh, dh], dtype=v.dtype, device=v.device)

        canvas = torch.zeros([batch_size, self.h_dim, w, w], dtype=v.dtype, device=v.device)
        if self.x_dim > 0: x_ds = self.downsample(x)
        else: x_ds = torch.zeros([batch_size, 0, dh, dh], dtype=x.dtype, device=x.device)
        r = r.repeat(1, 1, dh // r.shape[2], dh // r.shape[3])
        v = v.repeat(1, 1, dh // v.shape[2], dh // v.shape[3])

        for l in range(self.core_count):
            # Compute prior distribution by extracting mean and logstd
            p_mean, p_logstd = torch.chunk(self.prior(gen_state_h), 2, dim=1)
            prior_distr = D.Normal(p_mean, F.softplus(p_logstd))

            # Compute generator output
            z = prior_distr.sample()

            gen_input = torch.cat([v, r, z, x_ds], dim=1)

            if self.share:
                cell = self.generator
            else:
                cell = self.generator[l]

            if self.gru:
                gen_state_h = cell(gen_input, gen_state_h)
            else:
                gen_state_h, gen_state_c = cell(gen_input, gen_state_h, gen_state_c)

            if self.upscale_share:
                canvas += self.upsample(gen_state_h)
            else:
                canvas += self.upsample[l](gen_state_h)

        x_mean = self.observation(canvas)
        return torch.exp(x_mean) - 1

class VGG19(nn.Module):
    def __init__(self, requires_grad=False):
        super(VGG19, self).__init__()
        features = torchvision.models.vgg19(pretrained=True).features
        self.feature1 = torch.nn.Sequential()
        for i in range(0, 4):
            self.feature1.add_module(str(i), features[i])

        self.feature2 = torch.nn.Sequential()
        for i in range(4, 9):
            self.feature2.add_module(str(i), features[i])

        self.feature3 = torch.nn.Sequential()
        for i in range(9, 16):
            self.feature3.add_module(str(i), features[i])

        self.feature4 = torch.nn.Sequential()
        for i in range(16, 23):
            self.feature4.add_module(str(i), features[i])

        self.feature5 = torch.nn.Sequential()
        for i in range(23, 32):
            self.feature5.add_module(str(i), features[i])

        if not requires_grad:
            for param in self.parameters():
                param.requires_grad = False

    def forward(self, X):
        output = []
        X = F.interpolate(X, size=(224, 224), mode='bilinear')
        X = self.feature1(X)
        output.append(X)
        X = self.feature2(X)
        output.append(X)
        X = self.feature3(X)
        output.append(X)
        X = self.feature4(X)
        output.append(X)
        X = self.feature5(X)
        output.append(X)
        return output

class UnetBlock(nn.Module):
    def __init__(self, input_dim, output_dim):
        super(UnetBlock, self).__init__()
        self.model = nn.Sequential(
                nn.Conv2d(input_dim, output_dim, 3, padding=1),
                nn.ReLU(),
                nn.Conv2d(output_dim, output_dim, 3, padding=1),
                nn.ReLU()
            )

    def forward(self, x):
        return self.model(x)

class UnetGenerator(nn.Module):
    def __init__(self, settings, index):
        super(UnetGenerator, self).__init__()
        self.v_dim = settings.model.pose_dim
        self.w = settings.model.generators[index].render_size
        self.loss = settings.model.generators[index].settings.loss
        self.index = index
        self.representations = settings.model.generators[index].representations

        self.x_dim = 0
        for buffer in settings.model.generators[index].query_passes:
            self.x_dim += get_buffer_length(buffer)

        self.out_dim = 0
        for buffer in settings.model.generators[index].output_passes:
            self.out_dim += get_buffer_length(buffer)

        self.r_dim = 0
        for r in self.representations:
            self.r_dim += settings.model.representations[r].representation_dim

        # Set up architecture
        self.down_block1 = UnetBlock(self.x_dim + self.r_dim + self.v_dim, 128)
        self.pool1 = nn.MaxPool2d(2, stride=2) # 32x32
        self.down_block2 = UnetBlock(128 + self.r_dim + self.v_dim, 256)
        self.pool2 = nn.MaxPool2d(2, stride=2) # 16x16
        self.down_block3 = UnetBlock(256 + self.r_dim + self.v_dim, 512)

        self.pool3 = nn.MaxPool2d(2, stride=2) # 8x8
        self.down_block4 = UnetBlock(512 + self.r_dim + self.v_dim, 512)

        self.pool4 = nn.MaxPool2d(2, stride=2) # 4x4
        self.down_block5 = UnetBlock(512 + self.r_dim + self.v_dim, 512)

        self.pool5 = nn.MaxPool2d(2, stride=2) # 2x2
        self.down_block6 = UnetBlock(512 + self.r_dim + self.v_dim, 512)

        self.pool6 = nn.MaxPool2d(2, stride=2) # 1x1
        self.down_block7 = UnetBlock(512 + self.r_dim + self.v_dim, 512)

        self.upscale4 = nn.ConvTranspose2d(512, 512, 4, stride=2, padding=1) # 1x1
        self.up_block4 = UnetBlock(1024, 512)

        self.upscale5 = nn.ConvTranspose2d(512, 512, 4, stride=2, padding=1) # 2x2
        self.up_block5 = UnetBlock(1024, 512)

        self.upscale6 = nn.ConvTranspose2d(512, 512, 4, stride=2, padding=1) # 4x4
        self.up_block6 = UnetBlock(1024, 512)

        self.upscale1 = nn.ConvTranspose2d(512, 512, 4, stride=2, padding=1) # 8x8
        self.up_block1 = UnetBlock(1024, 512)
        self.upscale2 = nn.ConvTranspose2d(512, 256, 4, stride=2, padding=1) # 16x16
        self.up_block2 = UnetBlock(512, 256)
        self.upscale3 = nn.ConvTranspose2d(256, 128, 4, stride=2, padding=1) # 32x32
        self.up_block3 = UnetBlock(256, 128)

        self.final = nn.Conv2d(128, self.out_dim, 1)

        self.l1_loss = nn.L1Loss()
        self.end_relu = settings.model.generators[index].settings.end_relu

    def eval(self, x, o, v, p, r):

        batch_size, num_observations, pose_dim, *pose_dims = p.shape
        p = p.view((batch_size, num_observations * pose_dim, *pose_dims))

        w = x.shape[2]

        r1_shape = (batch_size, self.r_dim, w, w)
        r2_shape = (batch_size, self.r_dim, w // 2, w // 2)
        r3_shape = (batch_size, self.r_dim, w // 4, w // 4)
        r4_shape = (batch_size, self.r_dim, w // 8, w // 8)
        r5_shape = (batch_size, self.r_dim, w // 16, w // 16)
        r6_shape = (batch_size, self.r_dim, w // 32, w // 32)
        r7_shape = (batch_size, self.r_dim, w // 64, w // 64)

        v1_shape = (batch_size, self.v_dim, w, w)
        v2_shape = (batch_size, self.v_dim, w // 2, w // 2)
        v3_shape = (batch_size, self.v_dim, w // 4, w // 4)
        v4_shape = (batch_size, self.v_dim, w // 8, w // 8)
        v5_shape = (batch_size, self.v_dim, w // 16, w // 16)
        v6_shape = (batch_size, self.v_dim, w // 32, w // 32)
        v7_shape = (batch_size, self.v_dim, w // 64, w // 64)

        # Go down through network
        M = torch.cat([x, r.expand(*r1_shape), v.expand(*v1_shape)], dim=1) # 64x64
        X1 = self.down_block1(M)
        X = self.pool1(X1) # 32x32

        M = torch.cat([X, r.expand(*r2_shape), v.expand(*v2_shape)], dim=1)
        X2 = self.down_block2(M)
        X = self.pool2(X2) # 16x16

        M = torch.cat([X, r.expand(*r3_shape), v.expand(*v3_shape)], dim=1)
        X3 = self.down_block3(M)
        X = self.pool3(X3) # 8x8

        M = torch.cat([X, r.expand(*r4_shape), v.expand(*v4_shape)], dim=1)
        X4 = self.down_block4(M)
        X = self.pool4(X4) # 4x4

        M = torch.cat([X, r.expand(*r5_shape), v.expand(*v5_shape)], dim=1)
        X5 = self.down_block5(M)
        X = self.pool5(X5) # 2x2

        M = torch.cat([X, r.expand(*r6_shape), v.expand(*v6_shape)], dim=1)
        X6 = self.down_block6(M)
        X = self.pool6(X6) # 1x1

        X = torch.cat([X, r.expand(*r7_shape), v.expand(*v7_shape)], dim=1)
        X = self.down_block7(X)

        # Go upwards
        X = self.upscale4(X) # 2x2
        X = torch.cat([X, X6], dim=1)
        X = self.up_block4(X)

        X = self.upscale5(X) # 4x4
        X = torch.cat([X, X5], dim=1)
        X = self.up_block5(X)

        X = self.upscale6(X) # 8x8
        X = torch.cat([X, X4], dim=1)
        X = self.up_block6(X)

        X = self.upscale1(X) # 16x16
        X = torch.cat([X, X3], dim=1)
        X = self.up_block1(X)

        X = self.upscale2(X) # 32x32
        X = torch.cat([X, X2], dim=1)
        X = self.up_block2(X)

        X = self.upscale3(X) # 64x64
        X = torch.cat([X, X1], dim=1)
        X = self.up_block3(X)
        X = self.final(X)
        if self.end_relu:
            X = F.relu(X)

        return X

    def sample(self, x, o, v, p, r):
        X = self.eval(x, o, v, p, r)
        if "logl1" in self.loss or "logssim" in self.loss:
            X = torch.exp(X) - 1
        return X

    def forward(self, x, y, o, v, p, r, iteration, additional=None):
        Y = self.eval(x, o, v, p, r)
        losses = {}

        # Different loss functions, we used logl1 + logssim
        if "logl1" in self.loss:
            logl1 = 2 * torch.abs(Y - torch.log1p(y))
            losses["generator" + str(self.index) + "_logl1"] = logl1
        if "l1" in self.loss:
            l1 = 0.5 * torch.abs(Y - y)
            losses["generator" + str(self.index) + "_l1"] = l1
        if "logssim" in self.loss:
            logssim = 1.0 - pytorch_ssim.ssim(torch.log1p(y), Y)
            losses["generator" + str(self.index) + "_logssim"] = logssim
        if "ssim" in self.loss:
            ssim = 1.0 - pytorch_ssim.ssim(y, Y)
            losses["generator" + str(self.index) + "_ssim"] = ssim
        if "nll" in self.loss:
            sigma = max(0.1 + (2.0 - 0.1) * (1 - iteration / 2e4), 0.1)
            ll = -D.Normal(Y, sigma).log_prob(y)
            losses["generator" + str(self.index) + "_nll"] = ll

        if "logl1" in self.loss or "logssim" in self.loss:
            Y = torch.exp(Y) - 1

        output = {}
        output["generated"] = Y
        output["losses"] = losses

        return output

# Pixel generator model
class PixelCNN(nn.Module):
    def __init__(self, settings, index):
        super(PixelCNN, self).__init__()
        self.v_dim = settings.model.pose_dim
        self.w = settings.model.generators[index].render_size
        self.loss = settings.model.generators[index].settings.loss
        self.index = index
        self.representations = settings.model.generators[index].representations

        self.x_dim = 0
        for buffer in settings.model.generators[index].query_passes:
            self.x_dim += get_buffer_length(buffer)

        self.r_dim = 0
        for r in self.representations:
            self.r_dim += settings.model.representations[r].representation_dim

        self.out_dim = 0
        for buffer in settings.model.generators[index].output_passes:
            self.out_dim += get_buffer_length(buffer)

        # Create model
        self.num_layers = settings.model.generators[index].settings.layers
        self.num_hidden = settings.model.generators[index].settings.hidden_units

        self.propagate_buffers = settings.model.generators[index].settings.propagate_buffers
        self.propagate_representation = settings.model.generators[index].settings.propagate_representation
        self.propagate_viewpoint = settings.model.generators[index].settings.propagate_viewpoint

        self.input = nn.Conv2d(self.x_dim + self.r_dim + self.v_dim, self.num_hidden, 1, stride=1)

        self.prop_dim = 0
        if self.propagate_buffers:
            self.prop_dim += self.x_dim
        if self.propagate_representation:
            self.prop_dim += self.r_dim
        if self.propagate_viewpoint:
            self.prop_dim += self.v_dim

        self.layers = nn.ModuleList([])
        for _ in range(self.num_layers):
            self.layers.append(nn.Conv2d(self.num_hidden + self.prop_dim, self.num_hidden, 1, stride=1))

        self.final = nn.Conv2d(self.num_hidden, self.out_dim, 1, stride=1)
        self.activation = nn.LeakyReLU()
        self.l1_loss = nn.L1Loss()
        self.bce_loss = nn.BCELoss()
        self.end_relu = settings.model.generators[index].settings.end_relu
        if "vgg" in self.loss:
            self.vgg = VGG19()

    def eval(self, x, o, v, p, r):
        r_shape = (r.shape[0], r.shape[1], x.shape[2], x.shape[3])
        v_shape = (v.shape[0], v.shape[1], x.shape[2], x.shape[3])

        r = r.expand(*r_shape)
        v = v.expand(*v_shape)

        X = torch.cat([x, v, r], dim=1)
        X = self.activation(self.input(X))

        for i in range(len(self.layers)):
            if self.propagate_buffers:
                X = torch.cat([X, x], dim=1)
            if self.propagate_representation:
                X = torch.cat([X, r], dim=1)
            if self.propagate_viewpoint:
                X = torch.cat([X, v], dim=1)
            X = self.activation(self.layers[i](X))

        X = self.final(X)
        if self.end_relu:
            X = F.relu(X)
        if "bce" in self.loss:
            X = X.clamp(0, 1)

        return X

    def sample(self, x, o, v, p, r):
        X = self.eval(x, o, v, p, r)
        if "logl1" in self.loss or "logssim" in self.loss or "logbce" in self.loss or "loghuber" in self.loss or "lognll" in self.loss or "logl2" in self.loss:
            X = torch.exp(X) - 1
        return X

    def forward(self, x, y, o, v, p, r, iteration, additional=None):
        Y = self.eval(x, o, v, p, r)
        losses = {}

        # Different loss functions, we used logl1 + logssim
        if "logl1" in self.loss:
            logl1 = 2 * torch.abs(Y - torch.log1p(y))
            losses["generator" + str(self.index) + "_logl1"] = logl1
        if "l1" in self.loss:
            l1 = 2 * torch.abs(Y - y)
            losses["generator" + str(self.index) + "_l1"] = l1
        if "logl2" in self.loss:
            logl2 = 2 * torch.abs(Y - torch.log1p(y)) * torch.abs(Y - torch.log1p(y))
            losses["generator" + str(self.index) + "_l2"] = 250 * logl2
        if "logssim" in self.loss:
            logssim = 1.0 - pytorch_ssim.ssim(torch.log1p(y), Y)
            losses["generator" + str(self.index) + "_logssim"] = logssim
        if "ssim" in self.loss:
            ssim = 1.0 - pytorch_ssim.ssim(y, Y)
            losses["generator" + str(self.index) + "_ssim"] = ssim
        if "nll" in self.loss:
            sigma = max(0.1 + (2.0 - 0.1) * (1 - iteration / 2e5), 0.1)
            ll = -D.Normal(Y, sigma).log_prob(y)
            losses["generator" + str(self.index) + "_nll"] = 0.2 * ll
        if "bce" in self.loss:
            losses["generator" + str(self.index) + "_bce"] = 10 * self.bce_loss(Y, y)
        if "vgg" in self.loss:
            y_norm = torch.log1p(y).clamp(0, 1)
            Y_norm = Y.clamp(0, 1)
            target_features = self.vgg(y_norm)
            output_features = self.vgg(Y_norm)
            vgg_loss = 0
            for i in range(len(target_features)):
                vgg_loss += torch.mean((target_features[i] - output_features[i]) * (target_features[i] - output_features[i]), dim=[1, 2, 3], keepdim=True)
            losses["generator" + str(self.index) + "_vgg"] = 0.2 * vgg_loss

        if "logl1" in self.loss or "logssim" in self.loss or "logl2" in self.loss:
            Y = torch.exp(Y) - 1

        output = {}
        output["generated"] = Y
        output["losses"] = losses

        return output

# "Generator" that sums together all the input buffers
# Expects all of them to be same dimensions
# and this network has no loss
# This can be used to sum up the results of two generators such as direct and indirect outputs
# Helps with final visualization
class SumNet(nn.Module):
    def __init__(self, settings, index):
        super(SumNet, self).__init__()
        self.passes = settings.model.generators[index].query_passes

    def sample(self, x, o, v, p, r):
        buf_len = get_buffer_length(self.passes[0])
        Y = x[:, 0:buf_len, :, :]
        accum = buf_len

        for i in range(1, len(self.passes)):
            buf_len = get_buffer_length(self.passes[i])
            buffer = x[:, accum:(accum+buf_len)]
            accum += buf_len
            Y += buffer

        return Y

    def forward(self, x, y, o, v, p, r, iteration, additional=None):
        Y = self.sample(x, o, v, p, r)
        output = {}
        output["generated"] = Y
        output["losses"] = {}

        return output


#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@author: Xiaobo Yang
@contact: hal_42@zju.edu.cn
@software: PyCharm
@file: server.py
@time: 2019/4/27 7:26
@desc:
"""
from functools import wraps
from flask import Flask
from flask import request, jsonify, make_response
from solve_multi import Solve

app = Flask(__name__)

def allowCROD(func):
     def wrapper(*arg, **kwargs):
         resp = make_response(func(*arg, **kwargs))
         resp.headers['Access-Control-Allow-Origin'] = '*'
         resp.headers['Access-Control-Allow-Methods'] = 'PUT,GET,POST,DELETE'
         resp.headers['Access-Control-Allow-Headers'] = '*'
         return resp
     return wrapper

@app.route('/GetMove',methods=['GET', 'POST', 'OPTIONS'])
@allowCROD
def GetMove():
    resp = make_response()
    if request.method == 'GET':
        app.logger.info('Get an GET request')

        resp = make_response('You Sent an Get Request')
        return resp
    elif request.method == 'POST':
        app.logger.info('Get an Post request')

        req_data = request.get_json()
        app.logger.info(req_data)

        resp = Solve(req_data)
        # resp.headers['Access-Control-Allow-Origin'] = '*'
        return resp
    elif request.method == 'OPTIONS':
        app.logger.info('Get an OPTION request')
        return resp
    else:
        app.logger.error('Unrecognized Request')

        resp = make_response("Unrecognized Request")
        return resp

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=8000, debug=True)


from django import template

import matplotlib.pyplot as plt
import numpy as np
plt.rcParams.update({'font.size':14})

register = template.Library()

@register.filter
def conta_respostas(valor1):
    contador = 0
    for valor in valor1:
        contador +=1
    return contador

@register.filter
def porcentagem(total_respotas, resposta_cliente):
    total = 0
    for i in total_respotas:
        total += 1
    detrator = 0
    promotor = 0
    for c in resposta_cliente:
        if((resposta_cliente == 'Insatisfeito') or (resposta_cliente == 'Muito insatisfeito') or (resposta_cliente == 'Neutro')):
            detrator +=1
        else:
            promotor +=1

    porcent_detrator = (detrator * 100) / total
    porcent_promotor = (promotor * 100) / total

    return porcent_detrator, porcent_promotor


import time

from django.contrib.contenttypes.fields import GenericRelation
from django.core.validators import FileExtensionValidator
from django.db import models
from likes.models import Like
from mutagen.mp3 import MP3

def podcasts_directory_path(instance, filename):
    return f'podcasts/{instance.podcast_name}/image/{filename}'.format(filename=filename)

class Podcast(models.Model):
    podcast_name = models.CharField(
        max_length=200, verbose_name="Назва подкасту")
    theme = models.CharField(
        max_length=200, verbose_name="#Тематика", help_text="# вказувати не потрібно")
    description = models.TextField(max_length=450,
                                   verbose_name="Опис", help_text="до 450 символів")
    image = models.ImageField(
        upload_to=podcasts_directory_path, verbose_name="Зображення")
    is_published = models.BooleanField(default=True, verbose_name="Публікація")
    created_at = models.DateTimeField(
        auto_now_add=True, verbose_name="Дата створення")
    updated_at = models.DateTimeField(
        auto_now=True, verbose_name="Дата оновлення")

    def __str__(self):
        return self.podcast_name

    class Meta:
        ordering = ['id']
        verbose_name = "Подкаст"
        verbose_name_plural = "Подкасти"

def podcasts_audio_directory_path(instance, filename):
    return f'podcasts/{instance.podcast}/{instance.podcast_song_name}/{filename}'.format(filename=filename)

def podcasts_cover_directory_path(instance, filename):
    return f'podcasts/{instance.podcast}/{instance.podcast_song_name}/{filename}'.format(filename=filename)

class PodcastSong(models.Model):
    likes = GenericRelation(Like)
    podcast = models.ForeignKey(
        Podcast, blank=False, on_delete=models.CASCADE, verbose_name="Виконавець")
    release = models.IntegerField(default=0)
    podcast_song_name = models.CharField(max_length=200)
    description = models.TextField(
        max_length=450, blank=True, help_text="до 450 символів")
    cover = models.ImageField(upload_to=podcasts_cover_directory_path)
    audio = models.FileField(upload_to=podcasts_audio_directory_path, validators=[
                             FileExtensionValidator(['mp3'])])
    duration = models.CharField(
        max_length=10, blank=True, default='', verbose_name="Тривалість", help_text="Виглядає у форматі 00:00, вираховується автоматично")
    is_published = models.BooleanField(default=True)
    created_at = models.DateTimeField(auto_now_add=True)
    updated_at = models.DateTimeField(auto_now=True)

    def __str__(self):
        return self.podcast_song_name

    class Meta:
        ordering = ['-release']
        verbose_name = "Пісня подкасту"
        verbose_name_plural = "Пісні подкасту"

    def save(self, *args, **kwargs):
        if self.audio:
            audio = MP3(self.audio)
            ty_res = time.gmtime(audio.info.length)
            res = time.strftime("%M:%S", ty_res)
            self.duration = res
        super().save(*args, **kwargs)

    @property
    def total_likes(self):
        return self.likes.count()


import datetime

############################################################################################ printMonitorResults ######
def printMonitorResults(monitorResults):
    """  Prints the Monitor results
    """
    liveResults = monitorResults.totalMonitors-monitorResults.ScrappedMonitors-monitorResults.FailedMonitors
    liveInDate  = monitorResults.inDateMonitors+monitorResults.outDateMonitors-monitorResults.errDateMonitors

    inDate  = monitorResults.inDateMonitors  / liveInDate * 100
    outDate = monitorResults.outDateMonitors / liveInDate * 100

    print("-" + " Monitor Information " + "-"*58)
    print()
    print(f" Total Monitors    : {monitorResults.totalMonitors:4}             Monitors in date       : {monitorResults.inDateMonitors}  {inDate:.2f}%")
    print(f" Scrapped Monitors : {monitorResults.ScrappedMonitors:4}             Monitors out of date   : {monitorResults.outDateMonitors}  {outDate:.2f}%")
    print(f" Failed Monitors   : {monitorResults.FailedMonitors:4}             Monitors with errors   : {monitorResults.errDateMonitors}")
    print("="*80)
    print(f" Live Monitors     : {liveResults:4}                                      {liveInDate}")
    print()

############################################################################################ parseMonitors ############
def parseMonitors(monitors, monitorResults, logger):
    """  Loops through the directory monitors collating the Monitor results.
    """
    monitorResults.totalMonitors = len(monitors)

    for monitor in monitors:
        if monitor.Status == "Scrapped":
            monitorResults.ScrappedMonitors = monitorResults.ScrappedMonitors + 1
        elif monitor.Status == "Faulty":
            monitorResults.FailedMonitors = monitorResults.FailedMonitors + 1
        else:
            #  Only process live monitors.

            try:
                checkForErrors(monitor, monitorResults, logger)

                if monitor.PPMDueDate > datetime.datetime.now():
                    monitorResults.inDateMonitors = monitorResults.inDateMonitors + 1
                else:
                    monitorResults.outDateMonitors = monitorResults.outDateMonitors + 1
            except ValueError as err:
                logger.error(err.args)
                monitorResults.errDateMonitors = monitorResults.errDateMonitors + 1

######################################################################################### checkForErrors ######
def checkForErrors(monitor, monitorResults, logger):
    """  Check for some common errors with the data.
    """
    error_flag = False

    future_date = datetime.datetime.now() + datetime.timedelta(days=365)

    if monitor.PPMDueDate == None:                                      #  In case of error.
        raise ValueError(f"{monitor.serialNumber} has no PPM date set.")

    if monitor.monitorType == "":
        raise ValueError(f"Monitor has no monitor type {monitor.serialNumber}.")

    if monitor.Site == "":
        raise ValueError(f"Monitor has no Site {monitor.serialNumber}.")

    if monitor.Status not in ("Scrapped", "Faulty", None):
        raise ValueError(f"Monitor has unknown status {monitor.serialNumber} of {monitor.Status}.")

    if monitor.PPMDueDate > future_date:
        raise ValueError(f"Monitor has a future date {monitor.serialNumber} of {monitor.PPMDueDate}.")


#this is a program that records my invoices from sri
import sqlite3
from load_databases import load_databases

def main():
    #create a connection to db or create it

    connection = sqlite3.Connection("sri.db")

    #load or create databases
    load_databases(connection)

    #load gui

    #perform operations

if __name__ == "__main__":
    main()


from __future__ import absolute_import, division, print_function

import cmarkgfm

from .clean import clean

variants = {
    "GFM": cmarkgfm.github_flavored_markdown_to_html,
    "CommonMark": cmarkgfm.markdown_to_html,
}

def render(raw, variant="GFM", **kwargs):
    renderer = variants.get(variant)

    if renderer:
        rendered = renderer(raw)

        if rendered:
            return clean(rendered)

    return None


from image_generation import print_image
from config import death

def check_input(enter_choose):
    while not enter_choose.isdigit():
        print("...\nдевочки я в шоке\n\nПопросилже ЦИФРУ")
        enter_choose = input("> Напиши цифру варианта: ")
    return enter_choose
def wrong_death(string, chapter):
    print("— Ха-ха-ха, ты реально играешь не по правилам, таких не любим. Прикончу я тебя \n(я же просил напиши цифру)")
    print("Ты умер!")
    print("1. Переиграть главу\n2. Переиграть игру с нуля\n3. Выход")
    enter_choose = check_input(input("> Напиши цифру варианта: "))
    if enter_choose == "1":
        return chapter
    if enter_choose == "2":
        return 0
    if enter_choose == "3":
        return -1

def death_screen(chapter):
    print_image(death)
    print("Ты умер!")
    print("1. Переиграть главу\n2. Переиграть игру с нуля\n3. Выход")
    enter_choose = check_input(input("> Напиши цифру варианта: "))
    if enter_choose == "1":
        return chapter
    if enter_choose == "2":
        return 0
    if enter_choose == "3":
        return -1

def win_screen():
    print("Ты выйграл!")
    print("1. Переиграть игру с нуля\n2. Выход")
    enter_choose = check_input(input("> Напиши цифру варианта: "))
    if enter_choose == "1":
        return 0
    if enter_choose == "2":
        return -1

def next_page():
    item = input('Нажми Enter для перехода: ')
    while item != "":
        item = input('Да просто нажми ENTER')
    return


import urllib3 as ul3
from bs4 import BeautifulSoup
from selenium import webdriver
from selenium.webdriver.common.keys import Keys
from selenium.webdriver.common.action_chains import ActionChains
from selenium.webdriver.chrome.options import Options
from tqdm import tqdm
import sys
import copy
import time

from redeem2DB import FundDB

'''
    股票基金评分排名系统
    url : http://fund.eastmoney.com/data/fundranking.html
'''
#============= Configer =======================
bar_lengh = 60

fund_types = ['股票型', '混合型', '债券型', '指数型']
#============ Global Data =====================
data_array = [] # 基金数据 , shape = (array, map)
data_date = {'month' : 0, 'day' : 0} # 基金数据获取日期
fundDB = FundDB() #赎回连接对象
#========= Rubbish Data =========================
remove_data = [] #被删除的基金
#============ Tools ===========================
def web_driver():
    '''浏览器驱动'''
    chrome_options = Options()
    chrome_options.add_argument('--headless')
    # 无头浏览
    browser = webdriver.Chrome(chrome_options=chrome_options)
    # 有头浏览
##    browser = webdriver.Chrome()
    return browser
def init_Rank_Sort(map_):
    '''Add&Init key-[rank] to Map'''
    map_['rank'] = 0
def combine_rank(orig_, rank):
    '''combine rank from [rank] to [orig_]'''
    for index,rank_ in enumerate(rank):
        for orig__ in orig_:
            if orig__['基金代码'] == rank_['基金代码']:
                orig__['rank'] += index
def wait_action(browser, name=None):
    if name == None: name = ''
    else: name = '['+str(name)+']'
    print(name+'Waiting for Reflesh ...')
    while(BeautifulSoup(browser.page_source, "html.parser").get_text().find('长时间不更新') != -1):
        ''' Waiting for Reflesh '''
        pass
    print('Reflesh Done.')
#============ Get Fund Detail =======================
url = "http://fund.eastmoney.com/data/fundranking.html"
browser = web_driver()
print('Getting page ==> ', url)
browser.get(url)

for fund_type in fund_types:
    print('Getting Fund Data (', fund_type, ') ...')
    lis = browser.find_elements_by_tag_name('li')
    for li in lis:
        soup = BeautifulSoup(li.text, "html.parser")
        if soup.get_text().find(fund_type)>-1:
            ActionChains(browser).click(li).perform()
            break
    wait_action(browser)

    soup = BeautifulSoup(browser.page_source, "html.parser")
    trs = soup.find(id='dbtable').tbody.findAll('tr')
    if len(trs) <= 50:
        ActionChains(browser).click(browser.find_element_by_id('showall')).perform()
        wait_action(browser, '点击<showall>')
        soup = BeautifulSoup(browser.page_source, "html.parser")
        trs = soup.find(id='dbtable').tbody.findAll('tr')

    # Data Date
    try:
        tr = trs[0]
        tds = tr.findAll('td')
        date = tds[4].get_text().replace(' ', '').split('-')
        data_date['month'] = int(date[0])
        data_date['day'] = int(date[1])
    except ValueError:
        data_date['month'] = 0
        data_date['day'] = 0
    if data_date['month'] == 0 or data_date['day'] == 0:
        print('ERROR : 未能获取数据时间！')
    else:
        print('数据时间 (' + fund_type + ') -->', data_date['month'], '月', data_date['day'], '日')

    #Scope
    for tr in tqdm(trs, ncols=bar_lengh):
    ##for tr in trs[:20]:
        tds = tr.findAll('td')
        detail_map = {}
        detail_map['基金代码'] = tds[2].get_text().replace(' ', '')
        detail_map['基金简称'] = tds[3].get_text().replace(' ', '')
        try:
            detail_map['日增长率'] = float(tds[7].get_text().replace(' ', '').replace('%', ''))
            detail_map['近1周'] = float(tds[8].get_text().replace(' ', '').replace('%', ''))
            detail_map['近1月'] = float(tds[9].get_text().replace(' ', '').replace('%', ''))
            detail_map['近3月'] = float(tds[10].get_text().replace(' ', '').replace('%', ''))
        except ValueError:
            continue
        try:
            detail_map['近6月'] = float(tds[11].get_text().replace('---', '').replace(' ', '').replace('%', ''))
        except ValueError:
            detail_map['近6月'] = 0
        try:
            detail_map['近1年'] = float(tds[12].get_text().replace('---', '').replace(' ', '').replace('%', ''))
        except ValueError:
            detail_map['近1年'] = 0
        try:
            detail_map['近2年'] = float(tds[13].get_text().replace('---', '').replace(' ', '').replace('%', ''))
        except ValueError:
            detail_map['近2年'] = 0
        try:
            detail_map['近3年'] = float(tds[14].get_text().replace('---', '').replace(' ', '').replace('%', ''))
        except ValueError:
            detail_map['近3年'] = 0
        try:
            detail_map['今年来'] = float(tds[15].get_text().replace('---', '').replace(' ', '').replace('%', ''))
        except ValueError:
            detail_map['今年来'] = 0
        try:
            detail_map['成立来'] = float(tds[16].get_text().replace('---', '').replace(' ', '').replace('%', ''))
        except ValueError:
            detail_map['成立来'] = 0
        try:
            detail_map['手续费'] = float(tds[18].get_text().replace('---', '').replace(' ', '').replace('%', ''))
        except ValueError:
            detail_map['手续费'] = 100

        data_array.append(detail_map)

browser.quit()

#============== 加入赎回手续费 ===================================
url = "http://tools.fund.10jqka.com.cn/calculator/redeem.html"
browser = web_driver()
print('Getting page ==> ', url)
browser.get(url)
print('Adding Redemption Rates:')
#Scope
for map_ in tqdm(data_array, ncols=bar_lengh):
##for map_ in data_array:
    rate = None #赎回费率初始化
    rate = fundDB.get_redeem_by_code(map_['基金代码'])
    if rate == None:
        elem = browser.find_element_by_id('hqjij')
        elem.clear()
        elem.send_keys(map_['基金代码'])
        clock = time.clock()
        while True:
            if time.clock()-clock > 5 : break
            soup = BeautifulSoup(browser.page_source, "html.parser")
            if soup.find(id='hqjij').previous_sibling['style'] != 'left: 0px; top: 21px; width: 214px; z-index: 9999; display: none;':
                break

        promoteDis_1 = BeautifulSoup(browser.page_source, "html.parser").find(id='promoteDis')
        clock = time.clock()
        while True:
            if time.clock()-clock > 5 : break
            if promoteDis_1 != None:
                if BeautifulSoup(browser.page_source, "html.parser").find(id='promoteDis') != None \
                and BeautifulSoup(browser.page_source, "html.parser").find(id='promoteDis').get_text() != promoteDis_1.get_text():
                    break
            elif BeautifulSoup(browser.page_source, "html.parser").find(id='promoteDis') != None:
                break
            elem.send_keys(Keys.RETURN)

        soup = BeautifulSoup(browser.page_source, "html.parser")
        trs = soup.find(id='promoteDis').table.tbody.find_all('tr')
        trs = trs[1:]

        for tr in trs:
            try:
                rate_temp = float(tr.find_all('td')[1].get_text().replace('%','').replace(' ',''))
                if rate == None or rate < rate_temp:
                    rate = rate_temp
            except ValueError:
                pass

        if rate != None:
            fundDB.insert_redeem(map_['基金代码'], rate)

    if rate != None:
        map_['赎回费率'] = rate
    else:
        remove_data.append(copy.deepcopy(map_))
        del map_

fundDB.close()
browser.quit()

#============== Rank ============================================
for map_ in data_array:
    init_Rank_Sort(map_)

日增长率_sorted = sorted(data_array, key=lambda x : x['日增长率']-x['手续费']-x['赎回费率'], reverse=True)
近1周_sorted = sorted(data_array, key=lambda x : x['近1周']-x['手续费']-x['赎回费率'], reverse=True)
近1月_sorted = sorted(data_array, key=lambda x : x['近1月']-x['手续费']-x['赎回费率'], reverse=True)
近3月_sorted = sorted(data_array, key=lambda x : x['近3月']-x['手续费']-x['赎回费率'], reverse=True)
近6月_sorted = sorted(data_array, key=lambda x : x['近6月']-x['手续费']-x['赎回费率'], reverse=True)
近1年_sorted = sorted(data_array, key=lambda x : x['近1年']-x['手续费']-x['赎回费率'], reverse=True)
近2年_sorted = sorted(data_array, key=lambda x : x['近2年']-x['手续费']-x['赎回费率'], reverse=True)

combine_rank(data_array, 日增长率_sorted)
combine_rank(data_array, 近1周_sorted)
combine_rank(data_array, 近1月_sorted)
combine_rank(data_array, 近3月_sorted)
combine_rank(data_array, 近6月_sorted)
combine_rank(data_array, 近1年_sorted)
combine_rank(data_array, 近2年_sorted)

total_rank = sorted(data_array, key=lambda x : x['rank'])
print('=========================== Rank ======================================')
for index,item in enumerate(total_rank):
    print(index, item['基金代码'], item['基金简称'], item['rank'])

#=========== GUI for Search ======================
import tkinter as tk

root = tk.Tk()
root.title('股票基金评分排名')

def entry_Return(event):
    button_action()
entry = tk.Entry(root)
entry.bind('<Return>', entry_Return)
entry.grid(row=0, column=2)

def button_action():
    try:
        if int(entry.get()) >= len(total_rank):
            label['text'] = 'Sorry!No found.'
        else:
            detail = total_rank[int(entry.get())]
            label['text'] = 'rank:' + entry.get() + '\n'\
                            + '基金代码:' + str(detail['基金代码']) + '\n'\
                            + '基金简称:' + str(detail['基金简称']) + '\n'\
                            + '日增长率:' + str(detail['日增长率']) + '%\n'\
                            + '近1周:' + str(detail['近1周']) + '%\n'\
                            + '近1月:' + str(detail['近1月']) + '%\n'\
                            + '近3月:' + str(detail['近3月']) + '%\n'\
                            + '近6月:' + str(detail['近6月']) + '%\n'\
                            + '近1年:' + str(detail['近1年']) + '%\n'\
                            + '近2年:' + str(detail['近2年']) + '%\n'\
                            + '手续费:' + str(detail['手续费']) + '%\n'\
                            + '赎回费率:' + str(detail['赎回费率']) + '%'

    except ValueError:
        label['text'] = "Please input correct 'Number' !"
button = tk.Button(root, text='Check', command=button_action)
button.grid(row=0, column=3)

def button_code_action():
    code = ''
    try:
        int(entry.get())
        code = entry.get()
    except ValueError:
        code = 0
    detail = None
    rank_index = '---'
    for index,fund in enumerate(total_rank):
        if str(fund['基金代码']) == code:
            detail = fund
            rank_index = index
    if rank_index == '---':
        label['text'] = 'Sorry!No found.'
    else:
        label['text'] = 'rank:' + str(rank_index) + '\n'\
                        + '基金代码:' + str(detail['基金代码']) + '\n'\
                        + '基金简称:' + str(detail['基金简称']) + '\n'\
                        + '日增长率:' + str(detail['日增长率']) + '%\n'\
                        + '近1周:' + str(detail['近1周']) + '%\n'\
                        + '近1月:' + str(detail['近1月']) + '%\n'\
                        + '近3月:' + str(detail['近3月']) + '%\n'\
                        + '近6月:' + str(detail['近6月']) + '%\n'\
                        + '近1年:' + str(detail['近1年']) + '%\n'\
                        + '近2年:' + str(detail['近2年']) + '%\n'\
                        + '手续费:' + str(detail['手续费']) + '%\n'\
                        + '赎回费率:' + str(detail['赎回费率']) + '%'

button_code = tk.Button(root, text='Code', command=button_code_action)
button_code.grid(row=0, column=4)

label = tk.Label(root)
label.grid(row=1, column=2, columnspan=3)

root.mainloop()