#AML 6/8/18
#Prepare test data set- final
#Imports and PyRosestta init
#!/usr/bin/python

from pyrosetta import *
init()

from pyrosetta import toolbox
import rosetta.core.scoring.sasa as sasa
import rosetta.core.scoring.hbonds as hbonds
import numpy as numpy #for data output as a numpy array file, format is numpy.
from queue import PriorityQueue
from pyrosetta.rosetta.core.pose import compute_unique_chains
from math import sqrt, cos, sin
import numpy as np
import os
import sys
from multiprocessing import Pool, current_process
dataPath = os.getcwd()

#Creates N=15 clusters based on Calpha-Calpha distance from all residues in the protein
#Input: pyrosetta pose obejct of the protein
#Output: Dictionary w/ key of target residue pose res_id # and values of 15 closest neighbor pose res_id #
def getAllDistanceMatrix(pose):
	residues = {} #will be a dictionary of residue and fifteen closest neighbor residues
	#Will use PriorityQueeu class to get lowest 15
	for r in range(1,pose.total_residue() + 1):
		pri_q = PriorityQueue()
		for rsub in range(1,pose.total_residue()+1):
			if not rsub == r:
				ca_r = pose.residue(r).xyz('CA') #
				ca_rsub = pose.residue(rsub).xyz('CA')
				dist = sqrt((ca_r.x - ca_rsub.x)**2 + (ca_r.y - ca_rsub.y)**2 + (ca_r.z - ca_rsub.z)**2)
				#print (dist, " ", rsub)

				pri_q.put((dist, rsub))

		#get 15 closest neighbors

		residues[r] = []
		for i in range (1,16):
			stored = pri_q.get()
			st = list(stored)

			residues[r].append(st[1])
	return residues

#Creates a pyrosetta vector of all backbone SASA values for the residue (Ca, C, N, and O SASA values)
#Input: pose object of pro
#Output: a pyrosetta vector of the values
#TODO: ONLY RUN THIS ONCE
def getAllbbSASA(pose):
	calc = sasa.SasaCalc()
	calc.calculate(pose)
	return calc.get_residue_sasa_bb()

#get the number representation of Q3 secondary structure
#input: pose object of the protein, the pose res_id # for the residue you want the secondary structure for
#output: Whole number repres of the Q3 (1 =E, 2 = H, 3 = L)
def getQ3(pose, resid):
	char_ver = pose.conformation().secstruct(resid) #returns the secondary structure as L, H, or E
	char_to_numb = {'E':1, 'H':2, 'L':3}
	return char_to_numb[char_ver] #returns the number representation

#returns all hbonds in a dictionary, residue pose number key and list of donors and acceptors contents per key
#input of pose object for protein
#output of dictionary
#needs hbonds imported as hbonds from rosetta.core.scoring to work
#TODO: ONLY RUN ONCE
def getAllHBonds(pose):
    pose_hbonds = hbonds.HBondSet(pose, False) #gets all bb-bb only H-bonds for the pos
    hbond_dictionary = {}
    total_r = pose.total_residue()
    total_rplus = total_r + 1
    for r in range(1,total_rplus):
        hbond_dictionary[r] = []
        for h in range(1,pose_hbonds.nhbonds()+1):
            hb = pose_hbonds.hbond(h)
            a_res = hb.acc_res() #acceptor
            d_res = hb.don_res() #donor
            if r == a_res:
                hbond_dictionary[r].append(d_res)
            if r == d_res:
                hbond_dictionary[r].append(a_res)
    #
    return hbond_dictionary

#gets all three unit vector xyz portions
#input: pose.residue(#).xyz() for target residue (ca_pos2), and neighbor reisdue atoms (ca_pos, c_pos, n_pos)
#output: vector of all nine unit vector components [ca_cai, ca_caj, ca_cak, c_cai, c_caj, c_cak, n_cai, n_caj, n_cak]
#bool for passing, (true if all good, false if division by 0 occurs)
def vmu(ca_pos, c_pos, n_pos):
#####################  VECTOR   #####################
	vect = []
	failed = False
	#print ("FUNCTIONS!!!!!!!!!!!!!!!!!*******************!!!!!!!!!!!!!!!!!!!!")
	#pose.residue(neighbors[i]).xyz("CA"), pose.residue(neighbors[i]).xyz("C"), pose.residue(neighbors[i]).xyz("N"), pose.residue(r).xyz('CA'))
	#print (type(ca_pos))
	#print (ca_pos.item(0))
	vect.append(ca_pos.item(0))
	vect.append(ca_pos.item(1))
	vect.append(ca_pos.item(2))
	#print ("FUNNER FUNCTIONS!!!!!!!!!!!!!!!!!*******************!!!!!!!!!!!!!!!!!!!!")
	vect.append(c_pos.item(0)-ca_pos.item(0))
	vect.append(c_pos.item(1)-ca_pos.item(1))
	vect.append(c_pos.item(2)-ca_pos.item(2))

	vect.append(n_pos.item(0)-ca_pos.item(0))
	vect.append(n_pos.item(1)-ca_pos.item(1))
	vect.append(n_pos.item(2)-ca_pos.item(2))

	#print ("VECTOR: " , vect)

#####################  MAGNITUDE #####################
	magn = []
	magn1 = sqrt((vect[0]**2) + (vect[1]**2) + (vect[2]**2))
	magn2 = sqrt((vect[3]**2) + (vect[4]**2) + (vect[5]**2))
	magn3 = sqrt((vect[6]**2) + (vect[7]**2) + (vect[8]**2))
	if (magn1 == 0 or magn2 == 0 or magn3 == 0):
		failed = True #has failed
		uniVec = [0] *9
	else:
		magn.append(magn1)
		magn.append(magn2)
		magn.append(magn3)

		#print ("MAGNITUDE: ", magn)

	##################### UNIT VECTOR #####################
		uniVec = []
		uniVec.append(vect[0]/magn1)
		uniVec.append(vect[1]/magn1)
		uniVec.append(vect[2]/magn1)
		uniVec.append(vect[3]/magn2)
		uniVec.append(vect[4]/magn2)
		uniVec.append(vect[5]/magn2)
		uniVec.append(vect[6]/magn3)
		uniVec.append(vect[7]/magn3)
		uniVec.append(vect[8]/magn3)

	#print ("UNIT VECTOR: ", uniVec)
	return uniVec, failed

#x_output = open("./TrainingData/x_trainingdata.txt", "w")
#y_output = open("./TrainingData/y_trainingdata.txt", "w")
#testing out all data getters
def gen_binary_off_list(name_in):
	#with workers defaults names as a singl string, fixes by setting names as list containing incoming string
	j = 0
	names = [name_in]
	for name in names:
		print ("___________________________________________", name, "___________________________________________")
		try:
			pose = pose_from_pdb(name+".clean.pdb")
			total_len = pose.total_residue() #total protein length
			#total_datapoints = total_len * 293
			#print ("BIN STORED____________________________________________________")
			#bin_stored = np.fromfile("./binaryFiles/"+name+"binaryX", dtype = float, count = -1, sep = "")

			#need = True
			'''
			print ("ACTUAL: ", total_datapoints, " STORED: ", bin_stored.shape, " " ,bin_stored.shape[0])
			if not (bin_stored.shape[0] == total_datapoints): #so, failed the check
				print ("NEEDS FIXING")
				need = True
			else:
				print ("PASSED")
				need = False
			'''

			print (pose.pdb_info())
			dist_dict = getAllDistanceMatrix(pose)
			bb_list = getAllbbSASA(pose)
			hb_dict = getAllHBonds(pose)
			xarray = np.zeros((1,293))
			yarray = np.zeros((1,20))
			is_unique = compute_unique_chains(pose)
			failed_residue = False
			#print ("HELLO!")
			for r in range(1, pose.total_residue() + 1):
				if not is_unique[r]:
					continue
				x = []
				#print ("******************************")
				j = j + 1
				res = pose.residue(r)
				print ("On: ", r , " out of ", pose.total_residue())

				x.append(cos(pose.phi(r)))
				x.append(sin(pose.phi(r)))
				x.append(cos(pose.psi(r)))
				x.append(sin(pose.psi(r)))
				x.append(cos(pose.omega(r)))
				x.append(sin(pose.omega(r)))
				x.append(bb_list[r])
				x.append(getQ3(pose, r))

				neighbors = dist_dict[r] #list of neighbors
				ca_r = pose.residue(r).xyz('CA')
				ca_np = np.array([ca_r.x, ca_r.y, ca_r.z])
				c_r = pose.residue(r).xyz('C')
				c_np = np.array([c_r.x, c_r.y, c_r.z])
				n_r = pose.residue(r).xyz('N')
				n_np = np.array([n_r.x, n_r.y, n_r.z])
				nxp = np.subtract(n_np,ca_np) / np.linalg.norm(np.subtract(n_np,ca_np))
				cxp = np.subtract(c_np, ca_np)/np.linalg.norm(np.subtract(c_np, ca_np))
				xp = -1 *nxp #x axis
				zp = np.cross(cxp, xp)
				zp = zp / np.linalg.norm(zp)
				yp = np.cross(zp,xp)
				yp = yp / np.linalg.norm(yp)
				neg_ca_np = -1*ca_np
				rm = np.transpose(np.array([xp,yp,zp]))
				rm = np.vstack((rm, np.array([0, 0, 0])))
				rm = np.hstack((rm, [[0],[0],[0],[1]]))
				rm = np.matrix(rm)
				tm = np.matrix( [[1, 0, 0, 0], [0, 1, 0, 0],[0,0,1, 0], [-ca_r.x, -ca_r.y, -ca_r.z, 1]])
				m = tm * rm

				for i in range(0,15):
				#print ("NEIGHBORS!!!!!!!!!!!!!!!!!*******************!!!!!!!!!!!!!!!!!!!!")
					x.append(cos(pose.phi(neighbors[i])))
					x.append(sin(pose.phi(neighbors[i])))
					x.append(cos(pose.psi(neighbors[i])))
					x.append(sin(pose.psi(neighbors[i])))
					x.append(cos(pose.omega(neighbors[i])))
					x.append(sin(pose.omega(neighbors[i])))
					x.append(bb_list[neighbors[i]])
					ca_rsub = pose.residue(neighbors[i]).xyz('CA')
					x.append(sqrt((ca_r.x - ca_rsub.x)**2 + (ca_r.y - ca_rsub.y)**2 + (ca_r.z - ca_rsub.z)**2))
					rsub_ca_raw = pose.residue(neighbors[i]).xyz("CA")
					rsub_ca = np.array([rsub_ca_raw.x, rsub_ca_raw.y, rsub_ca_raw.z, 1])
					rsub_ca = np.array(rsub_ca *m)
					rsub_c = pose.residue(neighbors[i]).xyz("C")
					rsub_c = np.array([rsub_c.x, rsub_c.y, rsub_c.z, 1])
					rsub_c = np.array(rsub_c *m)
					rsub_n = pose.residue(neighbors[i]).xyz("N")
					rsub_n = np.array([rsub_n.x, rsub_n.y, rsub_n.z, 1])
					rsub_n = np.array(rsub_n *m)
					#print ("NEIGHBORS!!!!!!!!!!!!!!!!!*******************!!!!!!!!!!!!!!!!!!!!")
					holder, pass_vmu = vmu(rsub_ca, rsub_c, rsub_n)
					#print ("NEIGHBORS!!!!!!!!!!!!!!!!!*******************!!!!!!!!!!!!!!!!!!!!")
					#if there is not an error in the coordinates
					x.append(holder[0])
					x.append(holder[1])
					x.append(holder[2])
					x.append(holder[3])
					x.append(holder[4])
					x.append(holder[5])
					x.append(holder[6])
					x.append(holder[7])
					x.append(holder[8])
					x.append(getQ3(pose, neighbors[i]))
					x.append(hb_dict[r].count(neighbors[i]))
					if pass_vmu:
						print ("FAILURE DETECTED")
						failed_residue = True
				y = [0] * 20
				order_aa = ['A', 'R', 'N', 'D', 'C', 'Q', 'E', "G", "H", "I", "L", "K", "M", "F", "P", "S", "T", "W", "Y", "V"]
				seq = pose.sequence()
				#NOTE: pose.residue.aa() -> AA.aa_res for given residue identity
				#print (seq[r-1], " !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
				if seq[r-1] in order_aa:
					for i in range(0,20):
						if seq[r-1] == order_aa[i]:
							y[i] = 1
				else:
					failed_residue = True
				if not failed_residue:
					xarray = np.vstack((xarray,np.asarray(x)))
					yarray = np.vstack((yarray,np.asarray(y)))

			#use numpy vstack and asarray to add rows to a numpy array object
			print ("old array")
			print (xarray.shape)
			xarray = numpy.delete(xarray, (0), axis=0)
			yarray = numpy.delete(yarray, (0), axis=0)
			print ("fixed array")
			print (xarray.shape)

			if xarray.shape[0] > 0:
				xarray.tofile(name + "binaryX", sep ="")
				print ("GOOD______________________________________________________________________________________")
				yarray.tofile(name + "binaryY", sep="")
				#make five-fold cross val files
			else:
				failures.write(name)
				print ("BAD")
				failures.write("\n")

		except:
			print ("FAILURE AT: ", name)
			failures.write(name)
			failures.write("\n")

#running for single protein:
#1st-create a clean.pdb of the protein
#TODO: GET pro name, run this program
name = "2LZM"
toolbox.cleanATOM(name+".pdb")
#2nd-Run with gen_binary_off_list with single string name input
gen_binary_off_list(name)

'''
print("running program in folder "+ dataPath)
names_raw = open("misshapenFiles.txt", "r")
names = names_raw.read().splitlines()

failures = open("./failedResidues2.txt", "w")
#https://p16.praetorian.com/blog/multi-core-and-distributed-programming-in-python set up multiprocessing with a pool
print (sys.argv)
if __name__ == '__main__':
	with Pool(processes=8) as pool:
		#map doWork to availble Pool processes
		pool.map(gen_binary_off_list, names)

names_raw.close()
failures.close()
'''
#2nd Declare results variables and organization schemes
#x = [cos(phi) sin(phi) cos(psi) sin(psi) cos(omega) sin(omega) bbSASA Q3 15*(cos(phi) sin(phi) cos(psi) sin(psi) cos(omega) sin(omega) bbSASA CaCa_distance Ca_Ca_unitVector Ca_N_unitVector Ca_C_unitVector Q3 Number_Hydrogen_Bonds)]

#y = residue (1X20 vector, 1 one and rest 0's)

#x = [cos(phi) sin(phi) cos(psi) sin(psi) cos(omega) sin(omega) bbSASA Q3 number_residues_in_sphere 15*(cos(phi) sin(phi) cos(psi) sin(psi) cos(omega) sin(omega) bbSASA CaCa_distance Ca_Ca_unitVector Ca_N_unitVector Ca_C_unitVector Q3 Number_Hydrogen_Bonds hydrophobicity(from paper))]


class Solution:
  def minDifficulty(self, jobDifficulty: List[int], d: int) -> int:
    N = len(jobDifficulty)
    dp = [[float("inf") for x in range(d+1)] for y in range(N+1)] # some value for considering the first i jobs and having j cuts left

    def rec(cJob, cSum, cutsLeft):
      if cJob > N or cutsLeft < 0: return
      if dp[cJob][cutsLeft] <= cSum: return
      if cJob > 0:
        dp[cJob][cutsLeft] = min(dp[cJob][cutsLeft], cSum)

      cMax = 0
      for job in range(cJob+1, N+1):
        cMax = max(cMax, jobDifficulty[job-1])
        rec(job, cSum+cMax, cutsLeft-1) # do a cut at job

    rec(0, 0, d)

    return dp[N][0] if dp[N][0] != float("inf") else -1


formats = {
    'BMP': {
        'extension': '.bmp'
    },
    'JPEG': {
        'extension': '.jpg'
    },
    'PNG': {
        'extension': '.png'
    },
    'GIF': {
        'extension': '.gif'
    },
    'ICO': {
        'extension': '.ico'
    }
}

keys = list(formats.keys())

extensions = []
type_and_extensions = []

for key in keys:
    extensions.append(formats[key]['extension'])
    type_and_extensions.append((key, "*" + formats[key]['extension']))

# Tuple with all the extentions (All Formats)
extensions_string = ""
for extension in extensions:
    extensions_string += f"*{extension} "

extensions_string = extensions_string.strip()

all_formats = [('All Supporetd Formats', extensions_string)] + type_and_extensions


from KB_query.questions_model.classification import Base
from KB_query.questions_model.templates import *

# 平均比例问题查询条件生成器
class AvgProportion(Base):

    def __init__(self, word_objects, rules_list):
        super(AvgProportion, self).__init__(word_objects, rules_list)

        self.select = u"?y"
        self.select2 = u"?z"

        # 初始化参数
        self.__init_property()

    # 初始化参数
    def __init_property(self):
        self.init_time()

    # 生成器主函数
    def run(self):
        # Q18:1986年全国平均城镇人口占比是多少？
        if self.time is not None:
            return self.__time()

    # Q18:1986年全国平均城镇人口占比是多少？
    def __time(self):
        e = u"?x ex:UrbanPopulation ?y;\n" \
            u" ex:RuralPopulation  ?z;\n" \
            u" ex:year {time}." \
            .format(time=self.time.decode('utf-8'))

        return SPARQL_SUM_K \
            .format(prefix=SPARQL_PREXIX,
                    select=self.select,
                    select2=self.select2,
                    expression=e)


from common.bean.function_unit_status import FunctionUnitStatus
from common.bean.function_unit import FunctionUnit
from common.bean.instruction_status import InstructionStatus
from common.bean.instruction import Instruction
from common.bean.register_status import RegisterStatus
from common.bean.register_value import RegisterValue
from common.bean.stall import stall

class Frame:  # input file is in Frame list

    @staticmethod
    def newFrame(entries: dict):
        frame = Frame()

        frame.currentCycle = entries.get("currentCycle")

        frame.ProgramCounter = entries.get("ProgramCounter")

        frame.stallList = []
        stallList = entries.get("stallList")
        for ST in stallList:
            frame.stallList.append(stall.newStall(ST["type"], ST["dependToRegister"], ST["dependFromRegister"]))

        frame.instructionStatusList = []
        instructionStatusList = entries.get("instructionStatusList")
        for IS in instructionStatusList:
            frame.instructionStatusList.append(InstructionStatus.newInstructionStatus(
                Instruction.newInstruction(IS["instruction"]["issueCycle"], IS["instruction"]["tag"],
                                           IS["instruction"]["address"], IS["instruction"]["name"],
                                           IS["instruction"]["format"], IS["instruction"]["operandLeftName"],
                                           IS["instruction"]["operandRightName"], IS["instruction"]["destinationName"]),
                IS["stage"]))

        frame.functionUnitStatus = FunctionUnitStatus.newFunctionUnitStatus([])
        functionUnitList = entries.get("functionUnitStatus").get("functionUnitList")
        for FU in functionUnitList:
            frame.functionUnitStatus.functionUnitList.append(
                FunctionUnit.newFunctionUnit(FU["name"], FU["busy"], FU["op"], FU["des"], FU["fi"], FU["fj"], FU["fk"], FU["qj"], FU["qk"], FU["rj"], FU["rk"])
            )

        frame.registerStatusList = []
        registerStatusList = entries.get("registerStatusList")
        for RS in registerStatusList:
            frame.registerStatusList.append(RegisterStatus.newRegisterStatus(RS["registerName"], RS["functionUnitName"]))

        frame.registerValueList = []
        registerValueList = entries.get("registerValueList")
        for RV in registerValueList:
            frame.registerValueList.append(RegisterValue.newRegisterValue(RV["name"], RV["value"]))

        frame.log = entries.get("log")
        return frame

    currentCycle: int
    ProgramCounter: str
    stallList: list  # stalls happened in current cycle
    instructionStatusList: list  # InstructionStatus list  only updated instruction in current cycle
    functionUnitStatus: FunctionUnitStatus
    registerStatusList: list  # RegisterStatus list
    registerValueList: list  # RegisterValue list
    log: str


"""
This code is to scrape testing data for dozens of countries around the world compiled
by the good people of Wikipedia. The article this data figures in is
https://en.wikipedia.org/wiki/COVID-19_testing. The data is released under the Creative
Commons Attribution-Share-Alike License 3.0.
"""

import numpy as np
import requests

from bs4 import BeautifulSoup

from .utils import distribution

class Testing:
    url = "https://en.wikipedia.org/wiki/Template:COVID-19_testing"

    @property
    def countries(self) -> list:
        """
        Get countries from the Wikipedia table
        """
        return [
            a.string
            for a in self.table.find_all("a")
            if not a.string.startswith("[")
            and len(a.string) > 1
            and a.string != "testing"
        ]

    @property
    def tests(self) -> list:
        """
        Get numbers of tests administered from Wikipedia table
        """
        bd = self.table.find("tbody")
        rows = bd.find_all("tr")
        tests = []
        for i, r in enumerate(
            rows
        ):  # find_all isn't an iterator, so we have to do skip the header in this clunky way
            if i == 0:
                continue
            else:
                tests.append(r.find("td").string.replace("\n", ""))

        return tests

    @property
    def testing_data(self):
        return list(zip(self.countries, self.tests))

    def random_country_tests(self):
        """
        Returns number of tests for a random country
        """
        r = requests.get(self.url)
        self.soup = BeautifulSoup(r.content, features="html.parser")
        self.table = self.soup.find("table", {"class": "covid19-testing"})

        data = self.testing_data
        dist = distribution(len(data))
        idx = np.random.choice(len(data), size=1, p=dist)
        chosen = data[idx[0]]

        return {"country": chosen[0], "tests": chosen[1]}


import aylien_news_api
from aylien_news_api.rest import ApiException

# Configure API key authorization: app_id
aylien_news_api.configuration.api_key['X-AYLIEN-NewsAPI-Application-ID'] = '{{current_app_id}}'
# Configure API key authorization: app_key
aylien_news_api.configuration.api_key['X-AYLIEN-NewsAPI-Application-Key'] = '{{current_app_key}}'

# create an instance of the API class
api_instance = aylien_news_api.DefaultApi()

opts = {
  'title': 'trump',
  'sort_by': 'social_shares_count.facebook',
  'language': ['en'],
  'not_language': ['es', 'it'],
  'published_at_start': 'NOW-7DAYS',
  'published_at_end': 'NOW',
  'entities_body_links_dbpedia': [
    'http://dbpedia.org/resource/Donald_Trump',
    'http://dbpedia.org/resource/Hillary_Rodham_Clinton'
  ]
}

try:
    # List stories
    api_response = api_instance.list_stories(**opts)
    print("API called successfully. Returned data: ")
    print("========================================")
    for story in api_response.stories:
      print(story.title + " / " + story.source.name)
except ApiException as e:
    print("Exception when calling DefaultApi->list_stories: %s\n" % e)


class Solution:
    def rangeSum(self, nums: List[int], n: int, left: int, right: int) -> int:
        lt = []
        for i in range(n):
            s = 0
            for j in range(i, n):
                s += nums[j]
                lt.append(s)
        lt.sort()
        return sum(lt[left-1:right]) % 1000000007


import pytest
import os
import mock

from jenkins_head.HeadsManager import HeadsManager
from jenkins_head.HeadHandler import HeadHandler

class TestHeadsManager(object):

    @pytest.mark.parametrize('file2Test, elementCount', [
        ('inputFiles/one_head_config.yaml', 1),
        ('inputFiles/three_heads_config.yaml', 3),
        ('inputFiles/five_heads_config.yaml', 5)
    ])
    @mock.patch.object(HeadHandler, 'check')
    def test_HeadsManager(self, mock_check, file2Test, elementCount):
        headsManager = HeadsManager(self.__getFilePathRelative2ThisTest(file2Test))
        headsManager.checkAllHeads()

        assert mock_check.call_count == elementCount

    def __getFilePathRelative2ThisTest(self, filePath):
        return os.path.join(os.path.dirname(__file__), filePath)


import math
from multiprocessing import Pool

import numpy as np
from numpy import inf

from src.cluster.chromosome import Chromosome
from src.cluster.dm_chromosome import DMChromosome
from src.cluster.dm_population import DMPopulation
from src.cluster.population import Population
from src.common.fp import find_frequent_patterns
from src.common.utils import make_cluster_from_medoids

class ClusterSolver:
    def __init__(self, distances, Z=10, P=20, ng_arr=None, Pc=0.65, Pm=0.2, Pmb=0.05, k=None,
                 dm_size=4, dm_ng=5):
        self._distances = distances

        # Find k number based on number of customers K and
        # acceptable TSPTW subproblem size Z
        if k is None:
            K = len(distances)
            self._k = math.ceil(K / Z)
        else:
            self._k = k

        self._P = P

        if ng_arr is None:
            self._ng_arr = [25, ]
        else:
            self._ng_arr = ng_arr

        self._Pc = Pc
        self._Pm = Pm
        self._Pmb = Pmb

        self._numpy_rand = np.random.RandomState(42)

        self._dm_size = dm_size
        self._dm_ng = dm_ng
        self._dm_cur_priority_list = []

        self._dm_sup_threshold = 3
        self._dm_min_pattern_length = 3

        self._DM_MAX_PATTERNS_SEARCH = 2
        self._DM_MAX_PRIORITY_LIST_SIZE = self._k

    # параллелизм здесь внутри будет для режима data_mining
    def solve_cluster_core_data_mining(self):
        function_values = []

        # Заданное число итераций для dm режима
        global_best_res = None
        global_best_fitness = inf
        global_best_chrom = DMChromosome(self._k, self._distances, self._dm_cur_priority_list)
        for i in range(self._dm_ng):
            print("DM iteration: {}. Best fitness: {}".format(i, global_best_fitness))

            patterns = {}
            patterns_search_i = 0

            cur_best_res = None
            cur_best_fitness = None
            cur_best_chrom = None
            while patterns_search_i < self._DM_MAX_PATTERNS_SEARCH and len(patterns) == 0:
                print("DM. dm: {}, patterns_search_i: {}.".format(i, patterns_search_i))

                # Получаем список с кластерами для каждого запуска
                res, cur_best_res, cur_best_fitness, cur_best_chrom, cur_best_hist = self._make_multistart_genetic(i,
                                                                                global_best_chrom)

                function_values.extend(cur_best_hist)

                # Помещаем все кластеры из разных потоков в один глобальный список
                flatten_res = res.reshape(self._dm_size * self._k, res[0][0].shape[0])

                # запуск fp-growth для поиска лучшей хромосомы из набора res
                patterns = find_frequent_patterns(flatten_res.tolist(), self._dm_sup_threshold)
                # patterns, _ = FPGrowth.fpgrowth(flatten_res.tolist(), self._dm_sup_threshold)
                filtered_patterns = []

                for key in patterns:
                    if all(-1 != el and 0 != el for el in key):
                        filtered_patterns.append([patterns[key], list(key)])
                patterns = filtered_patterns

                # if patterns is not None:
                #     for sup, pattern in patterns:
                #         if all(-1 != el and 0 != el for el in pattern) and len(pattern) >= self._dm_min_pattern_length:
                #             filtered_patterns.append([sup, pattern])
                # patterns = filtered_patterns
                patterns.sort(key=lambda x: (x[0], len(x[1])), reverse=True)

                patterns_search_i += 1

            if len(patterns) != 0:
                for sup, pattern in patterns:
                    ind_to_replace = -1
                    for ind, item in enumerate(self._dm_cur_priority_list):
                        if pattern in item[1] and len(pattern) > len(item[1]):
                            ind_to_replace = ind
                            break

                    if ind_to_replace != -1:
                        self._dm_cur_priority_list[ind_to_replace] = pattern
                    else:
                        if len(self._dm_cur_priority_list) == self._DM_MAX_PRIORITY_LIST_SIZE:
                            self._dm_cur_priority_list[-1] = [sup, pattern]
                        else:
                            self._dm_cur_priority_list.append([sup, pattern])
                        self._dm_cur_priority_list.sort(key=lambda x: (x[0], len(x[1])), reverse=True)
            else:
                print("DM. Patterns not found.")

            if cur_best_fitness < global_best_fitness:
                global_best_res = cur_best_res
                global_best_fitness = cur_best_fitness
                global_best_chrom = cur_best_chrom

        return global_best_res, function_values

    def _make_multistart_genetic(self, dm_iter_ind, elite_chromosome):
        # Мультистарт
        rand_arr = self._numpy_rand.rand(self._dm_size, 1)
        np_rand_arr = [np.random.RandomState(int(100 * i)) for i in rand_arr]

        args = []
        for np_rand in np_rand_arr:
            dm_population = DMPopulation(self._P, self._k, self._distances, self._dm_cur_priority_list)
            current_best_chromosome, population = self._init(
                dm_population,
                elite_chromosome,
                np_rand
            )
            dm_population.chromosomes[0] = current_best_chromosome
            args.append((population, current_best_chromosome, np_rand, self._ng_arr[dm_iter_ind]))

        res = []
        cur_best_hist = []
        with Pool(self._dm_size) as p:
            result = p.starmap(self._solve, args)

            best_res = None
            best_chromosome = None
            cur_best_fitness = inf
            histories = []
            for chrom, history in result:
                clusters, _ = make_cluster_from_medoids(self._distances, self._dm_cur_priority_list, chrom.genes)
                res.append(clusters)

                if best_res is None:
                    best_res = clusters

                if chrom.fitness < cur_best_fitness:
                    cur_best_fitness = chrom.fitness
                    best_res = clusters
                    best_chromosome = chrom

                histories.append(history)

            histories = np.array(histories)
            cur_best_hist = [min([el.fitness for el in histories[:, i].tolist()]) for i in range(histories[0].size)]

        return np.array(res), best_res, cur_best_fitness, best_chromosome, cur_best_hist

    def solve(self):
        current_best_chromosome, population = self._init(Population(self._P, self._k, self._distances),
                                                         Chromosome(self._k, self._distances), self._numpy_rand)
        res_chromosome = self._solve(population, current_best_chromosome, self._numpy_rand, self._ng_arr[0])
        clusters, _ = make_cluster_from_medoids(self._distances, self._dm_cur_priority_list, res_chromosome.genes)
        return clusters

    def _solve(self, population, cur_best_chromosome, np_rand, ng):
        history = []

        self._print_current_iteration_info(0, cur_best_chromosome, population)
        history.append(cur_best_chromosome)
        for i in range(1, ng):
            population = population.selection(np_rand)
            population.crossover(self._Pc, self._Pmb, np_rand)
            population.mutate(self._Pm, np_rand)

            population.calculate_fitness()
            cur_best_chromosome = self._get_new_best_chromosome(cur_best_chromosome, population)

            self._print_current_iteration_info(i, cur_best_chromosome, population)
            history.append(cur_best_chromosome)

        return cur_best_chromosome, history

    def _init(self, population, chromosome, numpy_rand):
        population.generate_random_population(numpy_rand)
        population.calculate_fitness()
        current_best_chromosome = self._get_new_best_chromosome(chromosome, population)
        return current_best_chromosome, population

    def _print_current_iteration_info(self, i, current_best_chromosome, current_population):
        if i % 10 == 0:
            print("---------")
            print("Iteration: {}".format(i))
            print("Best chromosome fitness: {}".format(current_best_chromosome.fitness))
            # print("All chromosomes genes: {}".format(
            #     [chromosome.genes.tolist() for chromosome in current_population.chromosomes]))
            print("Best genes: {}".format(current_best_chromosome.genes))
            print("---------")

    def _get_new_best_chromosome(self, current_best_chromosome, population):
        best_chromosome_ind = population.find_best_chromosome_ind()
        best_chromosome = population.chromosomes[best_chromosome_ind]
        # print("cur fitness: {}".format(best_chromosome.fitness))
        if best_chromosome.fitness < current_best_chromosome.fitness:
            return best_chromosome
        return current_best_chromosome


import turtle
import colorsys
turtle.bgcolor("black")
turtle.tracer(100)
turtle.speed(100000)
def draw():
    h = 0
    n = 78
    turtle.up()
    turtle.goto(0, 60)
    turtle.down()
    turtle.pensize(5)
    for i in range(80):
        c = colorsys.hsv_to_rgb(h, 1, 1)
        h+=1/n
        turtle.color(c)
        turtle.fd(33)
        turtle.circle(i, 10, 3)
        for k in range(i):
            turtle.fd(25)
            turtle.rt(36)
            turtle.bk(9)
            turtle.circle(k, 5,11)
            turtle.rt(90.2)

draw()
for p in range(17):
    draw()
turtle.done()


import funcs as fn
import numpy as np
from scipy.optimize import fsolve
from scipy.integrate import odeint

def d_V_12(t_1=1.,t_2=1.,t_L=1,Re=1e8):
    """Function for reproducing Eqn. (16) in Ormel and Cuzzi (2007). """
    t_eta = Re**(-0.5)*t_L
    if (t_1 <= t_eta) and (t_2 <= t_eta):
        return np.sqrt(t_L / t_eta * (t_1 - t_2)**2)
    elif (t_1 > t_eta) and (t_1 < t_L):
        y_a = 1.6
        eps = t_1 / t_2
        return np.sqrt(2 * y_a - (1 + eps) + 2/(1 + eps) * (1/(1 + y_a) + eps**3/(y_a + eps))) * np.sqrt(t_1)
    elif (t_1 >= t_L):
        return np.sqrt(1/(1 + t_1) + 1/(1 + t_2))
    else:
        return np.sqrt(1/(1 + t_1) + 1/(1 + t_2))

def ts_sto_ram(r,params=[]):
    vth,mfp,rho_g,v_pg_tot,rho_s,tmp = params

    re = 4*r*v_pg_tot/vth/mfp
    cd = fn.drag_c(re)
    fd = 1./2.*cd*np.pi*r**2*rho_g*v_pg_tot**2

    t_s = 4./3.*np.pi*r**3*rho_s*v_pg_tot/fd

    return t_s

def st_solver(st=1,alph = 1e-100,a_au=1,m_suns=1,m_earths=1,verbose=0,smooth=1,gas_dep=1.,sig_g_in = 0, temp_in = 0):

    if verbose:
        print("St = %.7g" %st)
        print("a = %.3g AU" %a_au)
        print("m_star = %.3g m_sun" %m_suns)
        print("m_core = %.3g m_earths" %m_earths)

    a_core = fn.au*a_au
    m_star = fn.m_sun*m_suns
    r_core = fn.r_earth*((m_earths)**(1./3.))
    m_core = fn.m_earth*m_earths

    rho_obj = 2.0

    if sig_g_in:
        sig = sig_g_in
    else:
        sig = fn.surf_dens(a_core)

    if temp_in:
        t = temp_in
    else:
        t =  fn.temp(a_core)

    #Return derived parameters
    om = fn.omega(m_star,a_core)
    cs = fn.sound_speed(t)
    h = cs/om
    rho_g = fn.gas_density(sig,h)
    mfp = fn.mean_free_path(fn.mu,rho_g,fn.cross)
    vth = fn.therm_vel(cs)
    vkep = fn.vkep(m_star,a_core)
    eta = fn.eta(cs,vkep)

    h_r = fn.hill_rad(m_core,a_core,m_star)
    b_r = fn.bondi_rad(m_core,cs)

    if verbose:
        print("sig = %s" %sig)
        print("T = %s" %t)
        print("om = %s" %om)
        print("cs = %s" %cs)
        print("H = %s" %h)
        print("rho_g = %s" %rho_g)
        print("mfp = %s" %mfp)
        print("vth = %s" %vth)

    r_eps = st*rho_g/rho_obj*vth/om

    if verbose:
        print("r_eps = %.5g" %r_eps)

    if r_eps < 9.*mfp/4.:
        if verbose:
            print("Epstein Regime")
        return r_eps
    elif verbose:
        print("Not in Epstein")

    #Directly calculate v_rel from given St
    t_s = st/om
    if verbose:
        print("t_s = %.7g" %t_s)

    re_f = fn.re_f(alph,cs,h,vth,mfp)
    v_turb = np.sqrt(alph)*cs
    v_gas_tot = np.sqrt(v_turb**2 + (eta*vkep)**2)

    v_pg_L = eta*vkep*st*np.sqrt(4.+st**2.)/(1+st**2.)
    t_eddy = om**(-1)/np.sqrt(1+v_pg_L**2/v_turb**2)

    stl = t_s/t_eddy
    v_pg_turb = fn.v_pg(v_turb,stl,re_f)

    v_pg_tot = np.sqrt(v_pg_turb**2 + v_pg_L**2)

    if verbose:
        print("Re_t = %.7g" %re_f)
        print("t_eddy = %.7g" %(t_eddy*om))
        print("st_L = %.7g" %stl)
        print("v_turb = %.7g" %v_turb)
        print("v_gas_tot = %.7g" %v_gas_tot)
        print("v_pg_L = %.7g" %v_pg_L)
        print("v_pg_turb = %.7g" %v_pg_turb)
        print("v_pg_tot = %.7g" %v_pg_tot)

    def ts_zero(r,params):
        t_s = params[-1]
        return ts_sto_ram(r,params)-t_s

    params = [vth,mfp,rho_g,v_pg_tot,rho_obj,t_s]
    r_sto_ram = fsolve(ts_zero,9.*mfp/4.,args=params)[0]
    return r_sto_ram

class Disk:
    """ This class contains all of the properties of the protoplanetary disk. """
    def __init__(self, alpha=1e-100, a_core_au=1, m_star_solarmasses=1, gas_dep=1., sol_gas_ratio=0.01,
                 sig_p_in=0, sig_g_in=0, temp_in=0):

        self.m_star = m_star_solarmasses * fn.m_sun # Defines mass of star
        self.a_core = a_core_au * fn.au # Converts semi-major axis into cgs units

        ### Set disk parameters ###
        if sig_g_in: self.sig_gas = sig_g_in
        else: self.sig_gas = fn.surf_dens(self.a_core)/gas_dep # Calculates gas surface density
        if sig_p_in: self.sig_solid = sig_p_in # Calculates solid surface density
        else: self.sig_solid = self.sig_gas * sol_gas_ratio

        # Calculates temperature of disk
        if temp_in: self.T = temp_in
        else: self.T = fn.temp(self.a_core)

        self.alpha = alpha
        self.om = fn.omega(self.m_star, self.a_core) # Orbital frequency
        self.cs = fn.sound_speed(self.T) # Sound speed
        self.H = self.cs/self.om # Gas Scale height
        self.rho_gas = fn.gas_density(self.sig_gas, self.H) #Volumetric Density of gas
        self.mfp = fn.mean_free_path(fn.mu, self.rho_gas, fn.cross) # Mean free path
        self.v_th = fn.therm_vel(self.cs) # Thermal velocity
        self.v_kep = fn.vkep(self.m_star, self.a_core) # Keplerian velocity of core
        self.eta = fn.eta(self.cs, self.v_kep) # η, a measure of the local gas pressure support
        self.v_rel_i = self.v_kep # Sets initial velocity of core relative to gas

        # Turbulent gas parameters, see e.g. Ormel and Cuzzi (2007)
        self.re_f = fn.re_f(alpha, self.cs, self.H, self.v_th, self.mfp) # Reynolds number, MAYBE?
        self.v_gas_turb = (alpha)**(1/2)*self.cs # Turbulent velocity of gas
        self.v_gas_tot = np.sqrt(self.v_gas_turb**2 + (self.eta * self.v_kep)**2) # Total velocity of gas

class Core(Disk):
    """ This class contains all of the properties of an accreting core. """
    def __init__(self, a_core_au=1, m_core_earthmasses=1, alpha=1e-100, m_star_solarmasses=1, gas_dep=1.,
                 sol_gas_ratio=0.01, rho_core=2., sig_p_in=0, sig_g_in=0, temp_in=0, r_shear_off=0, extend_rh=1, alpha_z=0, h_mod=1):
        super().__init__(alpha, a_core_au, m_star_solarmasses, gas_dep, sol_gas_ratio, sig_p_in, sig_g_in, temp_in)
        self.r_shear_off = r_shear_off
        self.extend_rh = extend_rh
        self.alpha_z = alpha_z
        self.h_mod = h_mod

        self.m_core = m_core_earthmasses * fn.m_earth # Converts core mass into cgs units
        self.r_core = ((3 * self.m_core)/(4 * np.pi * rho_core))**(1/3) # Calculates radius of core

        ### Set core parameters ###
        self.r_hill = fn.hill_rad(self.m_core, self.a_core, self.m_star) # Hill radius
        self.r_bondi = fn.bondi_rad(self.m_core, self.cs) # Bondi radius
        self.r_bondi = max(self.r_bondi, self.r_core) # Prevents Bondi radius from shrinking below core radius
        self.rho_core = rho_core

    def t_stop(self, s, rho_obj=2.):
        """ Calculates stopping time, given particle size s. """
        if s > 9./4 * self.mfp: # Solve iteratively if we're in the fluid regime
            # Calculate terminal velocity and stopping time by iterating over force law. We've ignored angle here.
            delta = 1 # Used to check if we've converged
            v_i = self.v_rel_i
            while np.abs(delta) > .001:
                re = fn.rey(s, v_i, self.v_th, self.mfp) # Reynolds number
                m_obj = fn.obj_mass(s, rho_obj) # Mass of accreting objects
                dc = fn.drag_c(re) # Drag coefficient
                f_d = fn.stokes_ram(s, dc, self.rho_gas, v_i) # Drag force (ram pressure regime)
                t_s = m_obj*v_i/f_d # Stopping time
                v_r_lam = fn.v_new_r(self.eta, self.v_kep, t_s, self.om) # Radial component of laminar velocity (particle rel. to gas)
                v_phi_lam = fn.v_new_phi(self.eta, self.v_kep, t_s, self.om) # Phi component of laminar velocity
                v_lam = np.sqrt(v_r_lam**2 + v_phi_lam**2) # Total laminar velocity
                t_eddy = (self.om * (1 + (v_lam/self.v_gas_turb)**2)**.5)**-1 # Eddy crossing time
                stl = t_s/t_eddy # Stokes number
                if stl > 10:
                    v_new_turb = self.v_gas_turb * np.sqrt(1 - (1 + stl)**-1)
                else:
                    v_new_turb = fn.v_pg(self.v_gas_turb, stl, self.re_f)
                v_new = np.sqrt(v_lam**2 + v_new_turb**2) # Calculates total velocity, adding laminar and turbulent
                delta = (v_new - v_i)/v_i
                v_i = v_new
            t_s = t_s
        else: # Applies Epstein drag law if in diffuse regime
            t_s = fn.ts_eps(rho_obj, self.rho_gas, s, self.v_th)

        # Calculate laminar velocities from converged stopping time.
        v_r_lam = fn.v_new_r(self.eta, self.v_kep, t_s, self.om) # Radial component
        v_phi_lam = fn.v_new_phi(self.eta, self.v_kep, t_s, self.om) # Phi component
        v_lam = np.sqrt(v_r_lam**2 + v_phi_lam**2) # Total velocity

        t_eddy = (self.om * (1 + (v_lam/self.v_gas_turb)**2)**.5)**-1 # Eddy crossing time
        stl = t_s/t_eddy
        tau_s = t_s * self.om # Dimensionless stopping time

        if stl > 10: # Avoids using OC07 expressions for large stopping time
            v_obj_gas_turb = self.v_gas_turb * np.sqrt(1 - (1 + stl)**-1)
        else:
            v_obj_gas_turb = fn.v_pg(self.v_gas_turb, stl, self.re_f)
        v_obj_gas = np.sqrt(v_lam**2 + v_obj_gas_turb**2)
        return np.array([t_s, v_r_lam, v_phi_lam, v_lam, t_eddy, stl, tau_s, v_obj_gas_turb, v_obj_gas])

    def drag_force(self, s, vrel_i, rho_g, mfp, vth):
        re = fn.rey(s, vrel_i, vth, mfp)
        dc = fn.drag_c(re)
        if s > 9.*mfp/4.:
            fd = fn.stokes_ram(s,dc,rho_g,vrel_i)
        else:
            fd = fn.epstein_drag(s,rho_g,vrel_i,vth)
        return fd

    def r_wish(self, s, rho_obj=2.):
        """ Calculates wind-shearing radius, given particle size s. """
        v_cap = self.v_gas_tot # Set relevant velocity for orbit capture
        self.f_drag_obj = self.drag_force(s, v_cap, self.rho_gas, self.mfp, self.v_th)
        # Drag force on core
        self.f_drag_core = self.drag_force(self.r_core, self.v_core_gas, self.rho_gas, self.mfp, self.v_th)

        self.m_obj = (4./3.*np.pi*s**3)*rho_obj # Calculates mass of accreted object
        self.delta_a = np.abs(self.f_drag_obj/self.m_obj - self.f_drag_core/self.m_core) # Differential acceleration between core and object

        self.r_ws = fn.wish_radius(self.delta_a, self.m_core, self.m_obj) # Wind shearing radius
        return self.r_ws

    def r_shear(self, s, rho_obj=2.):
        """ Calculates shearing radius, given particle size s. """
        self.m_obj = (4./3.*np.pi*s**3)*rho_obj # Calculates mass of accreted object

        if self.r_shear_off: r_shear = 1e100 # Checks for flag
        elif s > 9./4 * self.mfp:
            def r_shear_solver(r_shear): # Function to pass to f_solve to determine r_shear
                v_rel = r_shear * self.om
                re = fn.rey(s, v_rel, self.v_th, self.mfp)
                dc = fn.drag_c(re)
                f_drag = fn.drag_force(s, v_rel, dc, self.rho_gas, self.mfp, self.v_th)
                return r_shear - np.sqrt(fn.G * self.m_core * self.m_obj/f_drag)

            # Solution for r_shear in the ram regime
            self.r_shear_ram = (fn.G * self.m_core * self.m_obj/
                                (.5 * .47 * self.rho_gas * np.pi * s**2. * self.om**2))**(1./4)
            self.r_shear_an = (3. * self.tau_s)**(1./3) * self.r_hill
            # Guess is the minimum of the analtyic solution in a linear regime and the solution in Ram
            r_shear_guess = min(self.r_shear_ram, self.r_shear_an)
            self.r_sh = fsolve(r_shear_solver, r_shear_guess)[0]
        else:
            self.r_sh = 3.**(1./3) * self.r_hill * self.tau_s**(1./3)
        return self.r_sh

    def r_accretion(self):
        """ Calculates accretion radius. """
        min_1 = np.minimum(self.r_ws, self.r_sh)
        self.r_stab = np.minimum(min_1, self.r_hill)

        self.r_atm = np.minimum(self.r_bondi, self.r_hill) # Calculates atmospheric radius as minimum of shearing and hill radii. Mickey currently has this coded as maximum, but should be minimum
        self.r_acc = np.maximum(self.r_atm, self.r_stab) # Calculates Accretion radius as maximum of bondi radius and stability radius
        self.r_acc = np.maximum(self.r_acc, self.r_core) # Make sure that the accretion radius doesn't get below the physical radius
        return self.r_acc

    def set_velocities(self, disp=0):
        """ Calculates a bunch of velocities, given particle size s. """
        self.v_cross = fn.vkep(self.m_core, self.r_acc) # Orbit velocity about core

        self.v_obj_phi = -self.eta * self.v_kep * (1/(1 + self.tau_s**2)) # Object velocity relative to Keplerian #np.abs(self.v_phi_lam + self.v_gas_lam)
        self.v_core_phi = -self.eta * self.v_kep * (1/(1 + self.tau_s_core**2)) # Core velocity relative to Keplerian #np.abs(self.v_phi_lam_core + self.v_gas_lam)

        self.v_lam_iner = np.sqrt((self.v_obj_phi - self.v_core_phi)**2 + self.v_r_lam**2) # Laminar, relative to inertial frame

        self.v_obj_core_lam = np.sqrt((self.v_obj_phi - self.v_core_phi)**2 + (self.v_r_lam - self.v_r_lam_core)**2) # Velocity of object relative to core
        self.v_obj_core_turb  = self.v_gas_turb * d_V_12(t_1=self.stl_core,t_2=self.stl,t_L=1,Re=self.re_f) # Use expresioons from Ormel Cuzzi 2007 to get relative turbulent particle velocity from stopping time of core #np.sqrt(self.v_gas_turb**2 - self.v_core_gas_turb**2)
        self.v_obj_core = np.sqrt(self.v_obj_core_lam**2 + self.v_obj_core_turb**2)

        self.v_shear = self.r_acc * self.om # Shear velocity

        self.v_inf = np.maximum(self.v_obj_core, self.v_shear) # Sets v_infinity
        self.v_enc = fn.G * self.m_core/self.r_acc/self.v_inf # Applies impulse approx to calculate encounter velocity

        if self.v_inf > self.v_cross: # Checks if impulse approximation is OK
            self.v_gas_enc = max(self.v_enc, self.v_obj_gas)
            self.v_grav = self.v_enc
        else:
            self.v_gas_enc = max(self.v_cross, self.v_obj_gas)
            self.v_grav = self.v_cross

    def encounter(self, s, rho_obj=2.):
        """ Calculate drag force, work, and kinetic energy during encounter, given particle size s.
            Also calculates the accretion probability. """
        self.re_enc = fn.rey(s, self.v_gas_enc, self.v_th, self.mfp) # Reynolds number during encounter
        self.dc_enc = fn.drag_c(self.re_enc) # Drag coefficient during encounter
        self.f_drag_enc = fn.drag_force(s, self.v_gas_enc, self.dc_enc, self.rho_gas, self.mfp, self.v_th) # Drag force
        self.work_enc = 2 * self.f_drag_enc * self.r_acc # Work done by drag over course of encounter
        self.ke = .5 * self.m_obj * self.v_inf**2 # Kinetic energy of object during encounter

        # Modify growth time by the ratio of the kinetic energy to work done over one orbit
        if (self.r_acc == self.r_hill and self.v_inf == self.v_shear and self.extend_rh):
            self.prob = min(self.work_enc/self.ke, 1.)
        else:
            self.prob = 1.

    def scale_heights(self):
        """ Calculate the turbulent scale height, Kelvin-Helmholtz scale height, disk scale height,
            accretion height, and accretion area. """
        if not(self.alpha_z): # Checks if vertical turbulence is different from other directions
            self.alpha_z = self.alpha
        self.H_turb = min(np.sqrt(self.alpha_z/self.tau_s) * self.H, self.H) # Turbulent scale height
        # Kelvin-Helmholtz scale height
        self.H_KH = fn.H_KH(self.H, self.a_core) * self.h_mod * min(1., 1/np.sqrt(self.tau_s))
        self.H_disk = np.maximum(self.H_turb, self.H_KH) # Scale height of disk, also called h_p

        self.H_acc = np.minimum(self.H_disk, self.r_acc) # Accretion height
        self.area_acc = 4 * self.r_acc * self.H_acc # Area over which objects are accreted

    def t_accretion(self):
        """ Calculate the accretion/growth time for an object of size s. """
        self.t_acc = (fn.growth_time(self.m_core, self.H_disk, self.sig_solid, self.area_acc, self.v_inf)
                      * fn.sec_to_years/self.prob)

        # Check energy criterion for accretion
        if (((self.r_stab > self.r_bondi and self.work_enc < self.ke) or (self.r_stab < self.r_bondi and
              self.work_enc > self.ke)) and not
           (self.r_stab == self.r_hill and self.v_inf == self.v_shear and self.extend_rh and self.work_enc < self.ke)):
            self.t_acc = 0 # Really should be infinity, but set to 0 for easy plotting
        return self.t_acc

    def main(self, s, rho_obj=2., disp=0):
        """ Runs each method defined for this class, in order to calculate and set all of the attributes
            of the object. """
        ts_arr = self.t_stop(s, rho_obj=rho_obj)
        self.t_s, self.v_r_lam, self.v_phi_lam, self.v_lam, self.t_eddy, self.stl, self.tau_s, self.v_obj_gas_turb, self.v_obj_gas = ts_arr
        ts_c_arr = self.t_stop(self.r_core, self.rho_core)
        self.t_s_core, self.v_r_lam_core, self.v_phi_lam_core, self.v_lam_core, self.t_eddy_core, self.stl_core, self.tau_s_core, self.v_core_gas_turb, self.v_core_gas = ts_c_arr
        r_ws = self.r_wish(s, rho_obj=rho_obj)
        r_sh = self.r_shear(s, rho_obj=rho_obj)
        r_acc = self.r_accretion()
        self.set_velocities(disp=disp)
        self.encounter(s, rho_obj=rho_obj)
        self.scale_heights()
        t_acc = self.t_accretion()


import plotly.graph_objects as _go
from scipy import stats as _stats
import pandas as _pd
import plotly.express as _px

def _get(inputDataFrame, plotConfigData):
    fig = _go.Figure()

    x = plotConfigData["params"]["x"]
    del plotConfigData["params"]["x"]

    showFit = False
    if "trendline" in plotConfigData["params"] and plotConfigData["params"]["trendline"] == "ols":
        showFit = True
        del plotConfigData["params"]["trendline"]

    groupParams = []
    # add color and facet entries to possible data groupings
    if "color" in plotConfigData["params"]:
        groupParams.append(plotConfigData["params"]["color"])
    if "facet_row" in plotConfigData["params"]:
        groupParams.append(
            plotConfigData["params"]["facet_row"])
    if "facet_col" in plotConfigData["params"]:
        groupParams.append(
            plotConfigData["params"]["facet_col"])
    groupParams = list(set(filter(None, groupParams)))

    def getProbPlot(groupData, grouped=True):
        res = _stats.probplot(groupData[x])
        if grouped:
            return _pd.Series({"Theoretical quantiles": res[0][0], "Percent": _stats.norm.cdf(res[0][0])*100.0, x: res[0][1], "Fit": res[0][0]*res[1][0] + res[1][1]})
        else:
            return _pd.DataFrame({"Theoretical quantiles": res[0][0], "Percent": _stats.norm.cdf(res[0][0])*100.0, x: res[0][1], "Fit": res[0][0]*res[1][0] + res[1][1]})

    if len(groupParams) > 0:
        plotData = inputDataFrame.groupby(groupParams).apply(
            lambda x: getProbPlot(x)).apply(_pd.Series.explode).reset_index()
    else:
        plotData = getProbPlot(inputDataFrame, False)

    # create the output plot

    plotConfigData["params"]["hover_data"] = ["Percent"]

    scatterPlot = _px.scatter(
        plotData, y="Theoretical quantiles", x=x,  **plotConfigData["params"])
    for d in scatterPlot.data:
        fig.add_trace(d)
    if showFit:
        for d in _px.line(plotData, y="Theoretical quantiles", x="Fit", **plotConfigData["params"]).data:
            d.showlegend = False
            fig.add_trace(d)
    fig.update_layout(scatterPlot.layout)
    fig.update_yaxes(
            tickmode='array',
            tickvals=[-3.090232306, -2.326347874, -1.281551566, -0.841621234, -0.524400513, -0.253347103,
                        0, 0.253347103, 0.524400513, 0.841621234, 1.281551566, 1.644853627, 2.326347874, 3.090232306],
            ticktext=["0.1%", "1%", "10%", "20%", "30%", "40%", "50%",
                        "60%", "70%", "80%", "90%", "95%", "99%", "99.9%"]
    )

    return fig


import re

from datetime import tzinfo

PLACEHOLDER_PATTERN = '\{(?P<kw>[a-zA-Z]*)(?::(?P<fmt>[^}]*))?\}'

def get_dump_filename(config, database, now):
    def repl(matches):
        kw = matches.group('kw')
        if kw == 'name':
            return config['name']
        elif kw == 'hostname':
            return config['hostname']
        elif kw == 'username':
            return config['username']
        elif kw == 'database':
            return database
        elif kw == 'now':
            date = now
        elif kw == 'utcnow':
            date = now.astimezone(tzinfo('UTC'))
        else:
            raise Exception('Invalid placeholder ' + kw +
                            ' in filename pattern')

        fmt = matches.group('fmt')
        if fmt is None:
            return fmt.isoformat()
        return date.strftime(fmt)

    filename = re.sub(PLACEHOLDER_PATTERN, repl, config['filename'])
    return config['directory'] + filename

def get_dump_glob(config):
    def repl(matches):
        kw = matches.group('kw')
        if kw == 'name':
            return config['name']
        elif kw == 'hostname':
            return config['hostname']
        elif kw == 'username':
            return config['username']
        elif kw == 'database':
            return '*'
        elif kw == 'now':
            return '*'
        elif kw == 'utcnow':
            return '*'
        else:
            raise Exception('Invalid placeholder ' + kw +
                            ' in filename pattern')
    filename = re.sub(PLACEHOLDER_PATTERN, repl, config['filename'])
    return config['directory'] + filename

def get_dump_matcher(config):
    def repl(matches):
        kw = matches.group('kw')
        if kw == 'name':
            return config['name']
        elif kw == 'hostname':
            return config['hostname']
        elif kw == 'username':
            return config['username']
        elif kw == 'database':
            if repl.db_kw_present:
                return '[a-zA-Z0-9_]+'
            else:
                repl.db_kw_present = True
                return '(?P<database>[a-zA-Z0-9_]+)'
        elif kw == 'now':
            return '[^/]+'
        elif kw == 'utcnow':
            return '[^/]+'
        else:
            raise Exception('Invalid placeholder ' + kw +
                            ' in filename pattern')
    repl.db_kw_present = False
    pattern = re.sub(PLACEHOLDER_PATTERN, repl, config['filename'])
    if not repl.db_kw_present:
        raise Exception('Database name must be present in filename pattern '
                        'to allow prune')
    return re.compile(config['directory'] + pattern)


import pandas as pd
import numpy as np
import torch
import os
import logging, sys
from torch.utils.data import Dataset, DataLoader

from stock_language import ts_to_string, BPE, tokenize

logging.basicConfig(stream=sys.stderr, level=logging.DEBUG)

# Dataset of ALL stocks from a particular index (ex: NASDAQ)
class StockDataset(Dataset):
    def __init__(self, data_dir, close_only=True, as_tensor=True):
        self.data_dir = data_dir
        self.paths = os.listdir(self.data_dir)
        self.close_only = close_only
        self.as_tensor = as_tensor

    def __len__(self):
        return len(self.paths)

    def __getitem__(self, idx):
        ticker_path = os.path.join(self.data_dir, self.paths[idx])
        data = pd.read_csv(ticker_path)
        if self.close_only:
            data = data.loc[:,'Close']
        if self.as_tensor:
            data = torch.tensor(data)
        else:
            data = data.to_numpy()
        return self.paths[idx].split('.')[0], data

class StockLanguageDataset(Dataset):
    def __init__(self, data_dir, preprocess=True, close_only=True, as_language=True, mini_test=False):
        self.data_dir = data_dir
        self.paths = os.listdir(self.data_dir)
        self.close_only = close_only
        self.as_language = as_language

        if mini_test:
            self.paths = self.paths[:min(5, len(self.paths))]

        if as_language:
            self.SL_strings = []
            for path in self.paths:
                ##
                logging.info(f'Processing path {path}')
                ##
                path = os.path.join(self.data_dir, path)
                data = pd.read_csv(path)

                if preprocess:
                    data = data.dropna()

                if self.close_only:
                    data = data.loc[:,'Close']
                else:
                    # operation not yet supported
                    ...

                data = ts_to_string(data)
                self.SL_strings.append(data)
        else:
            self.SL_strings = []

    def __len__(self):
        return len(self.SL_strings)

    def __getitem__(self, idx):
        return self.SL_strings[idx]

def basicSLDatasetDemo():
    nasdaq = './stock_market_data/nasdaq/csv'
    nasdaq_dataset = StockLanguageDataset(nasdaq, mini_test=True)
    nasdaq_dataloader = DataLoader(nasdaq_dataset, batch_size=1, shuffle=False)

    # need to unbatch
    SL_strings = [SL_string[0] for SL_string in nasdaq_dataloader]
    tokenizer = BPE(SL_strings, n=500)
    print(tokenizer)

def main():
    basicSLDatasetDemo()

if __name__ == '__main__':
    main()


# -*- coding: utf-8 -*-
# Time    : 2019/4/5 22:20
# Author  : Mifen
# Email   : 2952277346@qq.com
# Github  : https://github.com/Amd794

import pathlib
import sys

sys.path.append(pathlib.Path())
from BaseLsystem import Pic

rules = {
    'F': 'aFF-[b-F+F]+[c+F-F]',
    'X': 'aFF+[b+F]+[c-F]'
}
angle = 25
init_path = 'FX'
p = Pic(init_path, rules, angle)
apply_path = p.apply_rules()
p.draw_path(apply_path)


"""
0과 1로 이뤄진 문자열 s
s에 있는 모든 숫자를 전부 같게 하고 싶다
s에서 연속된 하나 이상의 숫자를 잡고 모두 뒤집는 것 -> 0을 1로 1을 0으로

0001100
"""

s = str(input())
split0 = list(filter(None, s.split("0")))
split1 = list(filter(None, s.split("1")))

lengh0 = len(split1) # 2
lengh1 = len(split0) # 1

if lengh0 > lengh1:
    print(lengh1)
elif lengh0 < lengh1:
    print(lengh0)
else:
    print(lengh0)


###############################################################################

import os
import sys
import numpy as np
from scipy.io import savemat

EPIpath=sys.argv[1]
print(EPIpath)

numTimePoints=int(sys.argv[2])
print(numTimePoints)

QCfile_name=os.environ['QCfile_name']
QCfile_name = ''.join([QCfile_name,'.log'])
print("QCfile_name: ",QCfile_name)
f=open(QCfile_name, "a+")

## FD
file_fd=''.join([EPIpath,'/motionRegressor_fd.txt'])
if os.path.exists(file_fd):
    fd = np.loadtxt(file_fd)
    print("file_fd dimensions: ",fd.ndim )
    if fd.ndim > 1:
        fd_scrub = np.sum(fd,axis=1)
    elif fd.ndim == 1:
        fd_scrub = fd
else:
    print("file_fd not generated; using dummy vector of zeros")
    f.write( "\n %s file not generated \n" % file_fd)
    fd_scrub = np.zeros(numTimePoints)

n_fd_outliers = np.count_nonzero(fd_scrub)
print("number of fd_outliers: ", n_fd_outliers)
f.write( "\n number of fd_outliers: %d \n" % n_fd_outliers)

# compute mean and std of motionMetric
file_fd=''.join([EPIpath,'/motionMetric_fd.txt'])
if os.path.exists(file_fd):
    fd = np.loadtxt(file_fd)
    fd_mean = np.mean(fd,axis=0)
    fd_std = np.std(fd,axis=0)
    fd_min = np.amin(fd,axis=0)
    fd_max = np.amax(fd,axis=0)
    f.write( "\n motionMatric_fd stats: \n ")
    f.write( "Mean = %f \n" % fd_mean)
    f.write( "Std = %f \n" % fd_std)
    f.write( "Min = %f \n" % fd_min)
    f.write( "Max = %f \n" % fd_max)

## DVARS
file_dvars=''.join([EPIpath,'/motionRegressor_dvars.txt'])
if os.path.exists(file_dvars):
    dvar = np.loadtxt(file_dvars)
    print("file_dvars dimensions: ",dvar.ndim )
    if dvar.ndim > 1:
        dvars_scrub = np.sum(dvar,axis=1)
    elif dvar.ndim == 1:
        dvars_scrub = dvar
else:
    print("file_dvars not generated; using dummy vector of zeros")
    dvars_scrub = np.zeros(numTimePoints)
n_dvars_outliers = np.count_nonzero(dvars_scrub)
print("number of dvars_outliers: ", n_dvars_outliers)
f.write( "\n number of dvars_outliers: %d \n" % n_dvars_outliers)

# compute mean and std of motionMetric
file_dvar=''.join([EPIpath,'/motionMetric_dvars.txt'])
if os.path.exists(file_dvar):
    dvar = np.loadtxt(file_dvar)
    dvar_mean = np.mean(dvar,axis=0)
    dvar_std = np.std(dvar,axis=0)
    dvar_min = np.amin(dvar,axis=0)
    dvar_max = np.amax(dvar,axis=0)
    f.write( "\n motionMatric_dvar stats: \n ")
    f.write( "Mean = %f \n" % dvar_mean)
    f.write( "Std = %f \n" % dvar_std)
    f.write( "Min = %f \n" % dvar_min)
    f.write( "Max = %f \n" % dvar_max)

scrub = np.add(fd_scrub,dvars_scrub)
good_vols = scrub == 0
good_vols = good_vols.astype(bool).astype(int)
print("number of good vols: ",np.count_nonzero(good_vols))
f.write( "\n number of good vols: %d \n" % np.count_nonzero(good_vols))
print("number of outliers to be scrubbed from dataset: ",np.count_nonzero(scrub))
f.write( "\n number of outliers to be scrubbed from dataset: %d \n" % np.count_nonzero(scrub))

f.close()

fname=''.join([EPIpath,'/scrubbing_goodvols.npz'])
np.savez(fname, good_vols=good_vols)

fname=''.join([EPIpath,'/scrubbing_goodvols.mat'])
print("savign MATLAB file ", fname)
mdic = {"good_vols": good_vols}
savemat(fname, mdic)


class Solution:
    def matrixReshape(self, mat, r, c):
        """
        Input: mat = [[1,2],[3,4]], r = 1, c = 4
        Output: [[1,2,3,4]]
        """
        row = len(mat)
        column = len(mat[0])
        if row*column != r*c:
            return mat
        else:
            total = []
            reshapeMatrix = []
            for i in range(row):
                total += mat[i]
            for i in range(r):
                reshapeMatrix.append(total[i*c:i*c+c])
            return reshapeMatrix

if __name__ == '__main__':
    mat = [[1, 2], [3, 4]]
    r = 4
    c = 1
    s = Solution()
    print(s.matrixReshape(mat, r, c))


from typing import Dict

import pandas as pd
from tqdm import tqdm

from python import DOCUMENT_ID, SENTENCE_IDX, TOKEN_IDX, PUBLISH_DATE, TOKEN, MENTION_ID, TOKEN_IDX_FROM, TOKEN_IDX_TO, \
    MENTION_TYPE, TIMEX_CORENLP_TO_OURS
from python.common_components import CORENLP
from python.common_components.corenlp import CoreNlp
from python.handwritten_baseline import POS, MENTION_TEXT, LEMMA, TIMEX_NORMALIZED
from python.handwritten_baseline.pipeline.data.base import BaselineDataProcessorStage, MODE_EXTEND, \
    MODE_INTERSECT, MODE_REPLACE
from python.handwritten_baseline.pipeline.data.base import Dataset
from python.handwritten_baseline.pipeline.data.processing import left_join_predictions, outer_join_predictions

class CoreNlpPosNerProcessorStage(BaselineDataProcessorStage):

    def __init__(self, pos, config, config_global, logger):
        super(CoreNlpPosNerProcessorStage, self).__init__(pos, config, config_global, logger)

        if self.mode == MODE_REPLACE:
            raise ValueError(f"Mode '{MODE_REPLACE}' is not supported.")

    @staticmethod
    def _parse_document(corenlp: CoreNlp,
                        doc_tokens: pd.DataFrame,
                        publish_date: pd.datetime):
        assert doc_tokens.index.names == [DOCUMENT_ID, SENTENCE_IDX, TOKEN_IDX]
        doc_id = doc_tokens.index.get_level_values(DOCUMENT_ID).unique()[0]

        tokens = []
        time_entities = []

        is_pretokenized = True

        # if present, pass the document publish date
        doc_date = None
        if pd.notna(publish_date):
            doc_date = publish_date.strftime("%Y-%m-%dT%H:%M:%S")

        for sent_idx, sent_tokens in doc_tokens.groupby(SENTENCE_IDX):
            properties = {}

            # The correct property to pass for the document date is "date", see https://stackoverflow.com/a/46290254.
            # When using any of the ner.docdate.xxx options, the CoreNLP server restarts for every request, which is terribly slow.
            if doc_date is not None:
                properties["date"] = doc_date
            if is_pretokenized:
                properties["tokenize.whitespace"] = True
                properties["tokenize.options"] = "ptb3Escaping=false"
                properties["ssplit.eolonly"] = True

            sentence_pretokenized = sent_tokens[TOKEN].str.cat(sep=" ")
            annotations = corenlp.parse_sentence(sentence_pretokenized, properties=properties)

            # if no "sentence" was recognized, keep the span as it was
            sentences = annotations.sentence
            if not sentences:
                raise ValueError

            # CoreNLP might tell us there is more than 1 sentence in the given string. We ignore that and trust our own
            # previously applied sentence tokenization. Therefore, tokens are counted across the corenlp sentence boundaries
            # here.
            sent_tokens = []
            token_idx = 0
            mentions = []
            sent_temporal_mentions = []

            for sentence in sentences:
                for token in sentence.token:
                    word, pos, lemma = token.word, token.pos, token.lemma
                    sent_tokens.append({TOKEN_IDX: token_idx, TOKEN: word, POS: pos, LEMMA: lemma})
                    token_idx += 1
                for mention in sentence.mentions:
                    _type, text = mention.entityType, mention.entityMentionText

                    row = {MENTION_TYPE: _type,
                           MENTION_TEXT: text,
                           TOKEN_IDX_FROM: int(mention.tokenStartInSentenceInclusive),
                           TOKEN_IDX_TO: int(mention.tokenEndInSentenceExclusive)}

                    # specific handling for temporal mentions
                    if _type in TIMEX_CORENLP_TO_OURS.keys():
                        # skip any non-timex mentions
                        if not mention.HasField("timex"):
                            continue

                        # update mention type and keep normalized TIMEX if we can get it
                        row[MENTION_TYPE] = TIMEX_CORENLP_TO_OURS[mention.timex.type]
                        if mention.HasField("normalizedNER"):
                            row[TIMEX_NORMALIZED] = mention.normalizedNER
                        sent_temporal_mentions.append(row)
                    else:
                        mentions.append(row)

            # add document id and sentence idx to each row
            for row in sent_tokens:
                row[DOCUMENT_ID] = doc_id
                row[SENTENCE_IDX] = sent_idx
            tokens += sent_tokens

            for row in sent_temporal_mentions:
                row[DOCUMENT_ID] = doc_id
                row[SENTENCE_IDX] = sent_idx
            time_entities += sent_temporal_mentions

        # combine tokens and mention info into new dataframes
        tokens = pd.DataFrame(tokens).set_index([DOCUMENT_ID, SENTENCE_IDX, TOKEN_IDX]).sort_index()

        # turn the index into [DOCUMENT_ID, MENTION_ID]
        if time_entities:
            time_entities_reindexed = []
            for doc_id, df in pd.DataFrame(time_entities).groupby(DOCUMENT_ID):
                temp = df.reset_index(drop=True).reset_index().rename(columns={"index": MENTION_ID})
                time_entities_reindexed.append(temp)
            time_entities = pd.concat(time_entities_reindexed).set_index([DOCUMENT_ID, MENTION_ID]).sort_index()
        else:
            time_entities = None

        return tokens, time_entities

    def _process_dataset(self,
                         dataset: Dataset,
                         live_objects: Dict) -> Dataset:
        corenlp = live_objects[CORENLP]

        # data is pre-tokenized, but run in through CoreNLP one more time for POS and named entities
        self.logger.info("Performing POS-tagging and NER")
        tokens = []
        time_entities = []

        for doc_id, doc_tokens in tqdm(dataset.tokens.groupby(level=[DOCUMENT_ID]),
                                                               desc="Parsing documents",
                                                               total=len(dataset.documents),
                                                               mininterval=10):
            # if present, pass the document publish date
            publish_date = pd.NaT
            if PUBLISH_DATE in dataset.documents.columns:
                publish_date = dataset.documents.loc[dataset.documents[DOCUMENT_ID] == doc_id, PUBLISH_DATE].item()

            doc_tokens, doc_time_entities = self._parse_document(corenlp, doc_tokens, publish_date)
            tokens.append(doc_tokens)
            if doc_time_entities is not None:
                time_entities.append(doc_time_entities)

        tokens = pd.concat(tokens).sort_index()
        time_entities = pd.concat(time_entities).sort_index()
        time_entities = time_entities.loc[time_entities[TIMEX_NORMALIZED].notna()]

        # add pos and lemma to each token in the dataset
        dataset.tokens = dataset.tokens.merge(tokens[[POS, LEMMA]], left_index=True, right_index=True)

        if dataset.mentions_time is None:
            dataset.mentions_time = time_entities
        else:
            # Common task for MODE_EXTEND and MODE_INTERSECT: perform left-join: keep all gold entity annotations,
            # but add information from CoreNLP output
            columns_keep_system = time_entities.columns.drop([TOKEN_IDX_FROM, TOKEN_IDX_TO, SENTENCE_IDX, MENTION_TYPE])
            dataset.mentions_time = left_join_predictions(dataset.mentions_time, time_entities, dataset.mentions_time.columns, columns_keep_system)

            if self.mode == MODE_INTERSECT:
                self.logger.info("Intersected new annotations with dataset from previous pipeline stages.")
            elif self.mode == MODE_EXTEND:
                self.logger.info("Extending set of temporal expressions in dataset with the ones found by CoreNLP")

                mentions_time_to_add = outer_join_predictions(time_entities, dataset).copy()
                dataset.mentions_time = pd.concat([dataset.mentions_time, mentions_time_to_add]).sort_index()

        # add mention text for later
        def get_mention_text_from_mention(row: pd.Series) -> str:
            return " ".join(dataset.tokens.loc[(row.name[0], row[SENTENCE_IDX], slice(row[TOKEN_IDX_FROM], row[TOKEN_IDX_TO] - 1)), TOKEN].values)
        dataset.mentions_time[MENTION_TEXT] = dataset.mentions_time.apply(get_mention_text_from_mention, axis=1)
        return dataset

component = CoreNlpPosNerProcessorStage


import torch
from torchvision.transforms import Normalize

# CIFAR10 dataset unnormalize as it comes out of the iter
def unnormalize_cifar10(normed):

    mean = torch.tensor([0.4914, 0.4822, 0.4465])
    std = torch.tensor([0.2023, 0.1994, 0.201])

    unnormalize = Normalize((-mean / std).tolist(), (1.0 / std).tolist())
    a = unnormalize(normed)
    a = a.transpose(0,1)
    a = a.transpose(1,2)
    a = a * 255
    b = a.clone().detach().requires_grad_(True).type(torch.uint8)

    return b

def unnormalize_femnist(normed):
    mean = torch.tensor([0.1307])
    std = torch.tensor([0.3081])

    unnormalize = Normalize((-mean / std).tolist(), (1.0 / std).tolist())
    a = unnormalize(normed)
    b = a.clone().detach().requires_grad_(True)
    return b


# -*- coding: utf-8 -*-
"""
Created on Tue May 10 21:15:53 2022

@author: ngloc
"""

from imageai.Detection import VideoObjectDetection

vid_obj_detect = VideoObjectDetection()

vid_obj_detect.setModelTypeAsYOLOv3()

vid_obj_detect.setModelPath(r"D:/UTN/CamIA/Test/3.object-det-with-yolo/yolo.h5")
vid_obj_detect.loadModel()

detected_vid_obj = vid_obj_detect.detectObjectsFromVideo(
    input_file_path =  r"D:/UTN/CamIA/Test/10sec2340x1080_ImagenRecortada.mp4",
    output_file_path = r"D:/UTN/CamIA/Test/3.object-det-with-yolo/24fpstest_10sec2340x1080_ImagenRecortada_procesado.mp4",
    frames_per_second=24,
    log_progress=True,
    return_detected_frame = True
)


import os
import cv2
import subprocess
import sys
import numpy

dirLeft  = "/Volumes/MySanDisk/mysdk/bvt_model/HumanSeg_mobile_v2_unet_8_hhl_2023-11-16_put_BackGround_before_crop"
dirMid   = "/Volumes/MySanDisk/mysdk/bvt_model/HumanSeg_mobile_v2_unet_8_hhl_2023-11-06"
dirRight = "/Volumes/MySanDisk/mysdk/bvt_model/HumanSeg_u2net_small_16_sigmoid_2023-10-09_DataClean1_Pass"
dirGT    = "/Volumes/MySanDisk/mysdk/bvt_model/val_3_mask_clip"

OUTPUT_NAME = "L_Low_Clean_2---M_Low_Clean_1---R_High_Pass"
OUTPUT_VIDEO_PATH = f"/Volumes/MySanDisk/mysdk/bvt_model/compare/{OUTPUT_NAME}_temp.mp4"
OUTPUT_VIDEO_PATH_RESIZE = f"/Volumes/MySanDisk/mysdk/bvt_model/compare/{OUTPUT_NAME}.mp4"
#OUTPUT_VIDEO_PATH = None

# if OUTPUT_VIDEO_PATH is not None and os.path.exists(OUTPUT_VIDEO_PATH):
#     raise Exception(f"{OUTPUT_VIDEO_PATH} exists")

video_file_list = [
    # "VID_1_out.mp4",
    # "VID_2_out.mp4",
    # "VID_3_out.mp4",
    # "VID_4_out.mp4",
    # "VID_5_out.mp4",
    # "back_chain_on_head_out.mp4",
    # "right_side_black_chain_out.mp4",
    # "long_hair_out.mp4",
    # "show_hand_out.mp4",
    # #"turn_around_out.mp4"

    "Chuang_out.mp4",
    "KeTing_out.mp4",
    "Office_Road_out.mp4",
    "Office_Seat_Angle_out.mp4",
    "Office_Seat_HalfBody_out.mp4",
    "Office_Seat_Postion_out.mp4",
    "YangTai_out.mp4",
]

def metric_mse(pred, true):
    #print(f"{pred.shape}, {true.shape}")  # 已经是 (1280, 720), (1280, 720)
    return ((pred - true) ** 2).mean() * 1e3

def iou_pytorch_numpy(preds, target):

    pred_cls   = preds > 0.5
    target_cls = target > 0.5

    intersection = numpy.sum(pred_cls & target_cls)
    union = numpy.sum(pred_cls | target_cls)

    iou = numpy.sum((intersection + 1e-6) / (union + 1e-6))

    return iou

def meaning(data):
    return sum(data) / len(data)

def put_text(np_img, text, postion):
    text = text
    org = postion
    fontFace = cv2.FONT_HERSHEY_SIMPLEX
    fontScale = 1
    color = (0, 255, 255)
    thickness = 2
    lineType = cv2.LINE_AA
    cv2.putText(np_img, text, org, fontFace, fontScale, color, thickness, lineType)

NUMBER_VIDEOS = len(video_file_list)
QUIT_KEY = ord('q') # 113
LEFT_KEY  = 2
RIGHT_KEY = 3
SPACE_KEY = 32

# current_index = 0
# end_flag = False
# output_video_writer = None

# while (not end_flag):

#     video_file = video_file_list[current_index]

#     # 打开源视频
#     left_video   = cv2.VideoCapture(os.path.join(dirLeft,  video_file) )
#     middle_video = cv2.VideoCapture(os.path.join(dirMid,   video_file) )
#     right_video  = cv2.VideoCapture(os.path.join(dirRight, video_file) )

#     # 获取视频帧的宽度和高度
#     frame_width  = int(left_video.get(cv2.CAP_PROP_FRAME_WIDTH))
#     frame_height = int(left_video.get(cv2.CAP_PROP_FRAME_HEIGHT))
#     frame_num    = int(left_video.get(cv2.CAP_PROP_FRAME_COUNT))

#     frame_width_m  = int(middle_video.get(cv2.CAP_PROP_FRAME_WIDTH))
#     frame_height_m = int(middle_video.get(cv2.CAP_PROP_FRAME_HEIGHT))
#     frame_num_m    = int(middle_video.get(cv2.CAP_PROP_FRAME_COUNT))

#     frame_width_r  = int(right_video.get(cv2.CAP_PROP_FRAME_WIDTH))
#     frame_height_r = int(right_video.get(cv2.CAP_PROP_FRAME_HEIGHT))
#     frame_num_r    = int(right_video.get(cv2.CAP_PROP_FRAME_COUNT))
#     print(f"file {video_file} ")
#     print(f"left_video    frame width:{frame_width}   height:{frame_height}     frame_num:{frame_num}")
#     print(f"middle_video  frame width:{frame_width_m}   height:{frame_height_m}   frame_num:{frame_num_m}")
#     print(f"right_video   frame width:{frame_width_r}   height:{frame_height_r}   frame_num:{frame_num_r}")

#     # 每个视频都清空重新计算
#     iou_list_l = []
#     mse_list_l = []
#     iou_list_m = []
#     mse_list_m = []
#     iou_list_r = []
#     mse_list_r = []

#     gt_mask_video    = None  # GT
#     left_mask_video  = None  # 左mask
#     mid_mask_video   = None  # 中mask
#     right_mask_video = None  # 右mask
#     gt_mask_name   = video_file.replace('_out', '' , 1)
#     if dirGT is not None and os.path.exists(os.path.join(dirGT, gt_mask_name)):
#         gt_mask_path = os.path.join(dirGT, gt_mask_name)
#         print(f"gt_mask_path: {gt_mask_path}")

#         mask_name = video_file.replace('_out', '_mask' , 1)
#         left_mask_path  = os.path.join(dirLeft,  mask_name)
#         mid_mask_path   = os.path.join(dirMid,   mask_name)
#         right_mask_path = os.path.join(dirRight, mask_name)
#         #print(f"left_mask_path: {left_mask_path}")

#         all_flag = True
#         all_flag = all_flag and os.path.exists(gt_mask_path)
#         all_flag = all_flag and os.path.exists(left_mask_path)
#         all_flag = all_flag and os.path.exists(mid_mask_path)
#         all_flag = all_flag and os.path.exists(right_mask_path)

#         if all_flag :
#             gt_mask_video    = cv2.VideoCapture(gt_mask_path)
#             left_mask_video  = cv2.VideoCapture(left_mask_path)
#             mid_mask_video   = cv2.VideoCapture(mid_mask_path)
#             right_mask_video = cv2.VideoCapture(right_mask_path)

#     while (True):
#         ret,  img  = left_video.read()
#         ret1, img1 = middle_video.read()
#         ret2, img2 = right_video.read()

#         # 读取mask(如果存在)
#         gt_mask    = None
#         left_mask  = None
#         mid_mask   = None
#         right_mask = None
#         ret_mask   = True
#         if gt_mask_video is not None:
#             mask_ret1, gt_mask = gt_mask_video.read()
#             mask_ret2, left_mask = left_mask_video.read()
#             mask_ret3, mid_mask = mid_mask_video.read()
#             mask_ret4, right_mask = right_mask_video.read()
#             ret_mask = mask_ret1 and mask_ret2 and mask_ret3 and  mask_ret4
#             if ret_mask:
#                 gt_mask    = gt_mask[:, :, 2] / 255.0
#                 left_mask  = left_mask[:, :, 2] / 255.0
#                 mid_mask   = mid_mask[:, :, 2] / 255.0
#                 right_mask = right_mask[:, :, 2] / 255.0 # 归一化 iou是计算>0.5 并且 训练代码计算mse的mask是在0~1

#         if not ret or not ret2 or not ret1 or not ret_mask:
#             print(f"end of video read ret={ret} ret1={ret1} ret2={ret2}")
#             # 下一个视频
#             current_index = current_index + 1
#             if (current_index >= NUMBER_VIDEOS):
#                 current_index = 0
#                 # 如果是保存到视频文件 全部文件遍历完就退出
#                 if OUTPUT_VIDEO_PATH is not None:
#                     end_flag = True
#             break

#         merged_img = cv2.hconcat([img, img1, img2])

#         if gt_mask_video is not None: # 计算每帧的mse iou和目前的平均 mse iou, 只判断一个通道
#             iou_l = iou_pytorch_numpy(gt_mask, left_mask)
#             iou_m = iou_pytorch_numpy(gt_mask, mid_mask)
#             iou_r = iou_pytorch_numpy(gt_mask, right_mask)
#             mse_l = metric_mse(gt_mask, left_mask)
#             mse_m = metric_mse(gt_mask, mid_mask)
#             mse_r = metric_mse(gt_mask, right_mask)
#             iou_list_l.append(iou_l)
#             iou_list_m.append(iou_m)
#             iou_list_r.append(iou_r)
#             mse_list_l.append(mse_l)
#             mse_list_m.append(mse_m)
#             mse_list_r.append(mse_r)
#             #print(f"iou:{iou_l:.2f},{iou_m:.2f},{iou_r:.2f};  mse:{mse_l:.2f},{mse_m:.2f},{mse_r:.2f}")
#             #print(f"mean iou:{meaning(iou_list_l):.2f},{meaning(iou_list_m):.2f},{meaning(iou_list_r):.2f}")
#             #print(f"mean mse:{meaning(mse_list_l):.2f},{meaning(mse_list_m):.2f},{meaning(mse_list_r):.2f}")
#             put_text(merged_img, f"cur   iou:{iou_l:.2f}, mse:{mse_l:.2f}",(0 , 50))
#             put_text(merged_img, f"mean iou:{meaning(iou_list_l):.2f}, mse:{meaning(mse_list_l):.2f}",(0 , 100)) # y轴需要偏移100
#             put_text(merged_img, f"cur   iou:{iou_m:.2f}, mse:{mse_m:.2f}",(img.shape[1] , 50))
#             put_text(merged_img, f"mean iou:{meaning(iou_list_m):.2f}, mse:{meaning(mse_list_m):.2f}",(img.shape[1] , 100))
#             put_text(merged_img, f"cur   iou:{iou_r:.2f}, mse:{mse_r:.2f}",(img.shape[1] * 2 , 50))
#             put_text(merged_img, f"mean iou:{meaning(iou_list_r):.2f}, mse:{meaning(mse_list_r):.2f}",(img.shape[1] * 2 , 100))
#             put_text(merged_img, f"{gt_mask_name}", (0 , 150))

#         if OUTPUT_VIDEO_PATH is not None:
#             if output_video_writer is None:
#                 fourcc = cv2.VideoWriter_fourcc(*'mp4v')
#                 print(f"out file size w:{merged_img.shape[1]} h:{merged_img.shape[0]}")
#                 output_video_writer = cv2.VideoWriter(OUTPUT_VIDEO_PATH, fourcc, 30.0, (merged_img.shape[1], merged_img.shape[0]))
#             output_video_writer.write(merged_img)
#         else:
#             cv2.imshow(OUTPUT_NAME, merged_img)
#             ######## key process begin ########
#             key = cv2.waitKey(1) & 0xFF
#             if (key != 255):
#                 print(f" key = {key}")

#             if (key == SPACE_KEY):
#                 while True:
#                     key = cv2.waitKey(1000) & 0xFF
#                     if (key == SPACE_KEY):
#                         break
#                     elif (key == QUIT_KEY):
#                         end_flag = True
#                         break
#             elif (key == QUIT_KEY):
#                 end_flag = True
#                 break
#             elif (key == LEFT_KEY):
#                 # 上一个视频
#                 current_index = current_index - 1
#                 if (current_index < 0):
#                     current_index = NUMBER_VIDEOS - 1
#                 break
#             elif (key == RIGHT_KEY):
#                 # 下一个视频
#                 current_index = current_index + 1
#                 if (current_index >= NUMBER_VIDEOS):
#                     current_index = 0
#                 break
#             ######## key process begin ########

#         # end 视频
#     left_video.release()
#     middle_video.release()
#     right_video.release()
#     # next 视频

# cv2.destroyAllWindows()

# if output_video_writer is not None: # 必须关闭文件 否则后面转换会出错 moov atom not found
#     output_video_writer.release()
# else:
#     sys.exit()

# 降低码率 视频文件大小

# 定义ffmpeg命令
command = ['ffmpeg', '-i', OUTPUT_VIDEO_PATH, '-b:v', '6M' , OUTPUT_VIDEO_PATH_RESIZE]
print(f"reduce video: {' '.join(command)}")

# 执行命令
result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

# 输出结果
print(f"stdout = {result.stdout.decode()}" )
print(f"stderr = {result.stderr.decode()}" )

# 执行成功, 删除大的视频文件
if result.returncode == 0:
    print(f"reduce ok")
    #os.remove(OUTPUT_VIDEO_PATH)
else:
    print(f"[ERROR] reduce fail")

print(f"reduce done")


from odoo import api, fields, models, _
import odoo.addons.decimal_precision as dp
from odoo.exceptions import ValidationError

class SaleOrder(models.Model):
    _inherit = "sale.order"

    @api.depends('order_line.price_total')
    def _amount_all(self):
        """
        Compute the total amounts of the SO.
        """
        for order in self:
            amount_untaxed = amount_tax = amount_discount = 0.0
            for line in order.order_line:
                amount_untaxed += line.price_subtotal
                amount_tax += line.price_tax
                amount_discount += (line.product_uom_qty * line.price_unit * line.discount) / 100
            order.update({
                'amount_untaxed': amount_untaxed,
                'amount_tax': amount_tax,
                'amount_discount': amount_discount,
                'amount_total': amount_untaxed + amount_tax,
            })

    discount_type = fields.Selection([('percent', 'نسبة'), ('amount', 'كمية')], string='نوع الخصم',
                                     readonly=True,
                                     states={'draft': [('readonly', False)], 'sent': [('readonly', False)]},
                                     default='percent')
    discount_rate = fields.Float(' معدل الخصم ', digits=dp.get_precision('الخصم'),
                                 readonly=True, states={'draft': [('readonly', False)], 'sent': [('readonly', False)]})
    amount_untaxed = fields.Monetary(string='Untaxed Amount', store=True, readonly=True, compute='_amount_all',
                                     track_visibility='always')
    amount_tax = fields.Monetary(string='Taxes', store=True, readonly=True, compute='_amount_all',
                                 track_visibility='always')
    amount_total = fields.Monetary(string='Total', store=True, readonly=True, compute='_amount_all',
                                   track_visibility='always')
    amount_discount = fields.Monetary(string='مبلغ الخصم', store=True, readonly=True, compute='_amount_all',
                                      digits=dp.get_precision('Account'), track_visibility='always')
    amount_before_discount = fields.Float(compute='get_amount_before_discount',string=" الاجمالى قبل الخصم ")

    @api.depends('amount_discount','amount_untaxed')
    def get_amount_before_discount(self):
        for rec in self:
            tot = rec.amount_untaxed + rec.amount_discount
            rec.amount_before_discount = tot

    @api.onchange('discount_type', 'discount_rate', 'order_line')
    def supply_rate(self):
        for order in self:
            if order.discount_type == 'percent':
                for line in order.order_line:
                    line.discount = order.discount_rate
            else:
                total = discount = 0.0
                for line in order.order_line:
                    total += round((line.product_uom_qty * line.price_unit))
                if order.discount_rate != 0:
                    discount = (order.discount_rate / total) * 100
                else:
                    discount = order.discount_rate
                for line in order.order_line:
                    line.discount = discount
                    # print(line.discount)
                    new_sub_price = (line.price_unit * (discount/100))
                    line.total_discount = line.price_unit - new_sub_price

    def _prepare_invoice(self, ):
        invoice_vals = super(SaleOrder, self)._prepare_invoice()
        invoice_vals.update({
            'discount_type': self.discount_type,
            'discount_rate': self.discount_rate,
        })
        return invoice_vals

    def button_dummy(self):

        self.supply_rate()
        return True

    # def write(self, vals):
    #     print(self)
    #     print(self.order_line)
    #     for line in self.order_line:
    #         line.write({
    #             'price_subtotal': line.total_discount
    #         })
    #     return super(SaleOrder, self).write(vals)
    # @api.onchange('discount_type', 'discount_rate')
    # def warning_msg(self):
    #     # settings_discount = self.env['ir.config_parameter'].get_param('group_discount_per_so_line')
    #     local_fields = self.env['sale.order.line'].fields_get()
    #     print(local_fields)
    #     # settings_discount = order.discount
    #     # print(settings_discount)
    #     # if not settings_discount:
    #     #     raise ValidationError(_("You have enable discount from configuration setting"))

class SaleOrderLine(models.Model):
    _inherit = "sale.order.line"

    discount = fields.Float(string='Discount (%)', digits=(16, 20), default=0.0)
    total_discount = fields.Float(string="Total Discount", default=0.0, store=True)


import time
import urllib
import spotipy
from spotipy.oauth2 import SpotifyOAuth
import credetials

#Configs
SPOTIPY_CLIENT_ID = credetials.SPOTIPY_CLIENT_ID
SPOTIPY_CLIENT_SECRET = credetials.SPOTIPY_CLIENT_SECRET
SPOTIPY_REDIRECT_URI = credetials.SPOTIPY_REDIRECT_URI
SCOPE = credetials.SCOPE
DEVICE_ID = credetials.DEVICE_ID
PLY_URI = credetials.PLY_URI

#Spotify definition of authentication
spfy = spotipy.Spotify(auth_manager=SpotifyOAuth(client_id=SPOTIPY_CLIENT_ID,
                     client_secret=SPOTIPY_CLIENT_SECRET, redirect_uri=SPOTIPY_REDIRECT_URI, scope=SCOPE))

def internet_on():
    try:
        urllib.request.urlopen('http://google.com')
        return True
    except:
        return False

start_time = '09:00:00'
end_time = '20:00:00'
start_time = time.strptime(start_time, '%H:%M:%S')
end_time = time.strptime(end_time, '%H:%M:%S')

# main loop
while True:
    if not internet_on():
        print("No connection")
    else:
        now = time.localtime()
        playing = spfy.current_user_playing_track()
        if playing is None or (playing.get('is_playing') == False and playing != None):
            spfy.start_playback(device_id=DEVICE_ID, context_uri=PLY_URI)
            print("Started to play")
        if start_time <= now <= end_time:
            spfy.volume(100, device_id=DEVICE_ID)
        else:
            spfy.volume(0, device_id=DEVICE_ID)
    time.sleep(300)


import pyrender
import numpy as np
import trimesh
from numba import jit, cuda
import cv2
import time
import matplotlib.pyplot as plt

from PIL import Image
#from RealismEmbedding.Harmonization import Harmonization
from Communication.python.client import Client
from HumanPose.demo import PoseEstimator
import pickle

@jit(nopython=True)
def compute_mask(rows, cols, im_color, im_depth, obj_depth, mask_, result_):
    for ro in range(rows):
        for co in range(cols):
            if obj_depth[ro][co] != 0:
                if obj_depth[ro][co] > im_depth[ro][co]:
                    result_[ro][co] = im_color[ro][co]
                    mask_[ro][co] = 0.0

                else:
                    result_[ro][co] = color[ro][co]
                    mask_[ro][co] = 255.0
    return result_, mask_
def is_inside_sphere(point,center, radius):
    diff = point-center
    return diff.T@diff < radius*radius

if __name__ == "__main__":
    client = Client("10.152.42.93", 5002)

    client.send("WP3")

    fuze_trimesh = trimesh.load('data/plane/textured.obj')

    scale = trimesh.transformations.scale_matrix(0.5)
    window_size = (640, 480)
    harmonize = Harmonization()
    estimator = PoseEstimator()
    last_hit = 0
    points = 0
    fourcc = cv2.VideoWriter_fourcc(*'MP4V')
    out = cv2.VideoWriter('output.mp4', fourcc, 5, (640, 480))
    while 1:
        # Receive data and unserialize
        start_time = time.time()
        serializedData = client.receive(decode=False)
        data = pickle.loads(serializedData)

        img_depth = data["depthMap"]
        img_depth = img_depth.astype(np.float32)
        img_depth = img_depth / 1000
        # img_depth = ((img_depth * 3.96) / 4.0) + 0.04

        img_color = data["rgbFrame"]
        #Render the Points achieved, can be done later or some other way, but I think this works fine
        cv2.putText(img_color, "Points: {}".format(points),(0,50), cv2.FONT_HERSHEY_TRIPLEX, 2, (255, 0, 0))
        #Estimate Human Poses
        human_poses = estimator.process_frame(img_color)
        # Get coordinate of right_wrist  == Keypoint 4
        wrist_coord = human_poses[0].keypoints[4]
        # Take a 3 x 3 Patch around that coordinate in the depth map
        hand_patch = img_depth[wrist_coord[0]-1:wrist_coord[0]+2, wrist_coord[1]-1:wrist_coord[1]+2]
        mean_depth = np.mean(hand_patch)
        wrist_coord_3d = np.array([wrist_coord[0], wrist_coord[1], mean_depth])

        camera_pose = np.eye(4)
        intrinsics = data["depthIntrincics"]

        object_pose = data["objectPose"]
        ## "Game" Logic
        if is_inside_sphere(wrist_coord_3d, object_pose[:3, 3], 0.1):
            if time.monotonic()- last_hit > 1:
                last_hit = time.monotonic()
                points += 1
        print(object_pose)
        camera = pyrender.camera.IntrinsicsCamera(intrinsics[0][0], intrinsics[1][1], intrinsics[0][2],
                                                  intrinsics[1][2],
                                                  znear=0.04, zfar=4.0, name=None)

        r = pyrender.OffscreenRenderer(window_size[0], window_size[1])
        scene = pyrender.Scene(bg_color=[0, 0, 0], ambient_light=[1.0, 1.0, 1.0])
        mesh = pyrender.Mesh.from_trimesh(fuze_trimesh, poses=object_pose)
        scene.add(mesh)

        scene.add(camera, pose=camera_pose)
        color, depth = r.render(scene)
        dim = depth.shape
        row = dim[0]
        col = dim[1]
        mask = np.zeros_like(depth)
        result = np.copy(img_color)

        indexes = depth > img_depth
        result[indexes] = img_color[indexes]
        mask[indexes] = 0.0
        result[~indexes] = color[~indexes]
        mask[~indexes] = 255.0
        z_indexes = depth == 0
        result[z_indexes] = img_color[z_indexes]

        # uncomment below line to enable harmonization
        # frame = harmonize.compute(Image.fromarray(mask), Image.fromarray(result))
        # comment below line to enable harmonization
        frame = result
        frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
        cv2.namedWindow('RS_DepthMap', cv2.WINDOW_AUTOSIZE)
        cv2.imshow('RS_DepthMap', frame)
        out.write(frame)

        key = cv2.waitKey(1)


import sys
class Solution(object):
    def maxProfit(self, prices):
        """
        :type prices: List[int]
        :rtype: int
        """
        profit=0
        min_price=10000
        for price in prices:
            min_price = min(min_price, price)

            profit = max(profit, price-min_price)

        if profit != 0:
            return profit
        else:
            return 0


#https://dmoj.ca/problem/ccc13s3
#CCC '13 S3 - Chances of Winning

def score(i):
    if i == len(game):
        for i in range(4):
            if i != t and points[i] >= points[t]:
                return 0
        return 1

    ways = 0
    a, b = game[i]

    points[a] += 3
    ways += score(i+1)
    points[a] -= 3

    points[b] += 3
    ways += score(i+1)
    points[b] -= 3

    points[a] += 1
    points[b] += 1
    ways += score(i+1)
    points[a] -= 1
    points[b] -= 1

    return ways

t = int(input()) - 1
g = int(input())

points = [0]*4
game = [(0,1), (0,2), (0,3), (1, 3), (1, 2), (2, 3)]

for _ in range(g):
    a,b, sa, sb = (int(x) for x in input().split())
    a -= 1
    b -= 1
    if sa > sb:
        points[a] += 3
    elif sb > sa:
        points[b] += 3
    else:
        points[a] += 1
        points[b] += 1
    game.remove((a,b))

print(score(0))


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

def load_agent_data(dir):
    return pickle.load(open(f"{dir}/sess_data.p", 'rb'))

agent_data = load_agent_data('/home/tzhang/Dropbox/Caltech/VisionLab/RL/Neuro/RL_Neuro/data/analysis/experiment_50/sess_0')

import pdb; pdb.set_trace()
state_action_log = agent_data.state_act_history_trials

def analyze_agent_novel_states_visited(state_action_log, truncate = 2000):
    all_agents_sameconfig = []
    for trial_data in state_action_log:
        one_agent_novel_states = []
        visited_states = set()
        for bout in trial_data: # each episode is one bout
            for state, _, _, in bout:
                visited_states.add(state)
                one_agent_novel_states.append(len(visited_states))
            # +1 at end because terminal state not counted
            one_agent_novel_states.append(one_agent_novel_states[-1])
        all_agents_sameconfig.append(one_agent_novel_states)

    if any(np.array([len(states) for states in all_agents_sameconfig]) < truncate):
        print(f'* Issue: at least one agent does not have enough data (<{truncate} nodes)')
        import pdb; pdb.set_trace()
    else:
        # trim each trial, and find variance / SD
        data_for_averaging = []
        for agent_data in all_agents_sameconfig:
            data_for_averaging.append(agent_data[:truncate])
        # compute variance
        var = np.std(data_for_averaging, axis = 0)
        avg = np.average(data_for_averaging, axis = 0)
        plt.figure()
        x = [i for i in range(truncate)]
        color = 'C0'
        plt.fill_between(x, avg - var, avg + var, facecolor = color, alpha = 0.3)
        plt.plot(avg, color = color)
        plt.savefig('/media/tzhang/Tony_WD_4TB_2/5_maze_v4/B_animals_reward/analysis/test2.png',
                    bbox_inches = 'tight', )

analyze_agent_novel_states_visited(state_action_log)


import discord
from discord.ext import commands
import requests

_carbon_url = "https://carbonnowsh.herokuapp.com/"

def code_to_url(code: str) -> str:
    return f"{_carbon_url}?&code={trim_url(encode_url(code))}"

class Carbon(commands.Cog):
    def __init__(self, bot):
        self.bot = bot

    @commands.command()
    async def carbonate(self, ctx, *, code):

        carbon_url = code_to_url(code)
        r = requests.get(carbon_url)

        b = BytesIO(r.content)
        await ctx.send(file=discord.File(fp=b, filename="code.png"))

async def setup(bot):
    await bot.add_cog(Carbon(bot))