#!/usr/bin/python3

# ## {R2D2919B742E} ##
# ###########################################################################
# What if magic existed?
# What if a place existed where your every thought and dream come to life.
# There is only one catch: it has to be written down.
# Such a place exists, it is called programming.
#    - Scott Taylor Reph, RightThumb.com
# ###########################################################################
# ## {C3P0D40fAe8B} ##

##################################################
import sys, time
##################################################
import _rightThumb._construct as __
appDBA=__.clearFocus(__name__,__file__);__.appReg=appDBA;
def focus(parentApp='',childApp='',reg=True):
	global appDBA;f=__.appName(appDBA,parentApp,childApp);
	if reg:__.appReg=f;
	return f
import _rightThumb._base3 as _
fieldSet=_.l.vars(focus(),__name__,__file__,appDBA)
_.load()
##################################################
_v = __.imp('_rightThumb._vars')
_str = __.imp('_rightThumb._string')
##################################################

def sw():
	pass
	#b)--> examples
	_.switches.register( 'Notes', '-n,-note,-notes' )
	_.switches.register( 'Download', '-dl,-download' )
	_.switches.register( 'Search-Notes-Content', '-db,--c,-cnt,-content' )
	_.switches.register( 'Print', '-p,-print' )
	#e)--> examples
	# _.switches.register( 'Files', '-f,-fi,-file,-files','file.txt', isData='glob,name,data,clean', description='Files', isRequired=False )

# __.setting('require-list',['Files,Plus','File,Has']) # todo
# __.setting('require-list',['Pipe','Files'])
__.setting('receipt-log')
__.setting('receipt-file')
__.setting('myFileLocations-skip-validation',False)
__.setting('require-pipe',False)
__.setting('require-pipe||file',False)
__.setting('pre-error',False)
__.setting('switch-raw',[])

_.appInfo[focus()] = {
	# 'app': '8facG-jo0Cxk',
	'file': 'pyJoplin.py',
	'liveAppName': __.thisApp( __file__ ),
	'description': 'Query the LIVE Joplin Desktop Database',
		# _.ail(1,'subject')+
		# _.aib('one')+
	'categories': [
						'DEFAULT',
				],
	'usage': [
						# 'epy another',
						# 'e nmap',
						# '',
	],
	'relatedapps': [
						# 'p another -file file.txt',
						# '',
	],
	'prerequisite': [
						# 'p another -file file.txt',
						# '',
	],
	'examples': [
						_.hp('p pyJoplin  + alexandria -content'),
						_.linePrint(label='simple',p=0),
						'',
	],
	'columns': [
					# { 'name': 'name', 'abbreviation': 'n' },
					# { 'name': '{1}', 'abbreviation': '{0}', 'sort': '{2}' },
	],
	'aliases': [
					# 'this',
					# 'app',
	],
	'notes': [
					# {},
	],
}

_.appData[focus()] = {
		'start': __.startTime,
		'uuid': '',
		'audit': [],
		'pipe': False,
		'data': {
					'field': {'sent': [], 'received': [] }, # { 'label': '', 'context': [],  }
					'table': {'sent': [], 'received': [] },
		},
	}

def triggers():
	_.switches.trigger( 'Files', _.myFileLocations, vs=True )
	_.switches.trigger( 'Ago', _.timeAgo )
	_.switches.trigger( 'Folder', _.myFolderLocations )
	_.switches.trigger( 'URL', _.urlTrigger )
	_.switches.trigger( 'Duration', _.timeFuture )

_.l.conf('clean-pipe',True)
_.l.sw.register( triggers, sw )

########################################################################################
#b)--> examples
#d)--> code hints to quickly get started
	#n)--> inline examples
		# any(ele in 'scott5' for ele in list('0123456789'))
		# if _.switches.isActive('Test'): test(); return None;
		# result=[]; result=[ _.pr(line) for i, line, bi in _.numerate( _.isData(r=0) )]
		# bk=[];[  bk.append(rec['backup']) for rec in backupLog if path == rec['file']]; bk=bk[-1];
		# a=(1 if True else 0) <--#
		#!)--> m=[[row[i] for row in matrix] for i in range(4)]

	#n)--> python globals
		# for k in globals(): print(k, eval(k) )

	#n)--> webpage from url
		# for subject in _.caseUnspecific( line, needle ): line = line.replace( subject, _.colorThis( subject, 'green', p=0 ) )

	#n)--> webpage from url
		# requests=__.imp('requests.post')
		#!)--> data=str(requests.post(url,data={}).content,'iso-8859-1')

	#n)--> import and backup example
		# _bk = _.regImp( __.appReg, 'fileBackup' ); _bk.switch( 'Silent' ); _bk.switch( 'isRunOnce' ); _bk.switch( 'Flag', 'APP' ); _bk.switch( 'DoNotSchedule' )
		# _bk.switch( 'Input', path ); bkfi = _bk.action();

	#n)--> inline
		# for rel in [ subject for subject in _.isData(r=0) if _.showLine(subject) ]: print(rel)

	#n)--> banner
		# banner=_.Banner(app); goss=banner.goss;
#e)--> examples
########################################################################################
#n)--> start

def jpage(p): global config; global _type_; return config[_type_].url.replace('[page]',str(p))

def all_data(t='notes'):
	global config
	global index
	global _type_
	_type_ = t
	# _.saveYML(index,'pyJoplin.yml')
	if not _type_ in index:
		index[_type_] = {}
		index[_type_]['parents'] = {}
		index[_type_]['ids'] = {}
		index[_type_]['rev'] = {}
		RECORDS = []
		while not config[_type_].current == config[_type_].max:

			config[_type_].current+=1
			print(_type_, 'page:',config[_type_].current)
			# print(jpage(config[_type_].current))
			page = _.URL( jpage(config[_type_].current) )
			if page == '{"items":[],"has_more":false}':
				config[_type_].max=config[_type_].current
				# _.pr(page,c='red')
			else:
				data = simplejson.loads(page)
				# try:
				# index.table = index.table + data
				for rec in data['items']:
					RECORDS.append(rec)

		# for rec in _.tables.returnSorted( 'RECORDS', 'd.epoch', RECORDS ):
		for rec in _.tables.returnSorted( 'RECORDS', 'title', RECORDS ):
			# print(rec)
			# print(rec['title'])
			if rec['parent_id']:
				if not rec['parent_id'] in index[_type_]['parents']: index[_type_]['parents'][rec['parent_id']] = []
				index[_type_]['parents'][rec['parent_id']].append(rec['id'])
				index[_type_]['ids'][rec['id']] = rec
			else:
				index[_type_]['ids'][rec['id']] = rec
				index['top'].append(rec['id'])
				print('parent:',rec['id'])
		_.saveYML(index,'pyJoplin.yml')
				# except:
				#     pass
				#     print('error')

spent={}

def varDepth(paths):
	var = "['xXx']"
	return var.replace('xXx',"']['".join(paths))

base=''
def _hierarchy(paths):
	global index
	def _hier_(pa):
		global hierarchy
		global base
		xXx = _archy_(pa)
		ex = 'if not "'+xXx[0]+'" in hierarchy'+xXx[1]+': hierarchy'+varDepth(pa)+'={};'
		# print(ex)
		exec(ex)
		# hierarchy
		# print(xXx)
	def _archy_(pa):
		p=pa.copy()
		global hierarchy
		global base
		vd = varDepth(pa)
		ix=p[-1]
		p.pop()
		if not p:
			vx=base
			if not pa[0] in hierarchy:
				# print(base)
				hierarchy[pa[0]]={}
		else:
			vx=varDepth(p)
			# ix=p[-1]
		return ix, vx

	global hierarchy
	# print('bof')
	# print(paths)
	for i,ID in enumerate(paths):
		_.v.IDs[ID]=1
		if _.v.print: _.pr('_.v.IDs[ID]=1',ID,index['folders']['ids'][ID]['title'].strip())
		p=[]
		for ii,iD in enumerate(paths):
			p.append(iD)
			if i==ii: break
		_hier_(p)
	# _hier_(paths)

		# print(vx, ix)
	# print('eof')

def fullpath2(ID):
	paths = fullpath(ID)
	if not ID in paths: paths.append(ID)
	for p in paths: _.v.IDs[p]=1
	names=[]
	for p in paths:
		if p in index['folders']['ids']: names.append(index['folders']['ids'][p]['title'])
		if p in index['notes']['ids']:   names.append(index['notes']['ids'][p]['title'])
	return '  / '.join(names)

def fullpath(ID,r=0):
	global index
	if not r: _.v.path=[]
	_.v.path = [ID] + _.v.path
	# .append(ID)
	try:
		if ID in index['folders']['ids']:
			rec = index['folders']['ids'][ID]
			if rec['parent_id']:
				return fullpath(rec['parent_id'],r=1)
	except Exception as e:
		pass
		# _.e(e)
	if not ID in _.v.path: _.v.path.append(ID)
	for i in _.v.path: _.v.IDs[i]=1
	_hierarchy(_.v.path)
	PATHS=[]
	for p in _.v.path:
		if not p in PATHS: PATHS.append(p)
	return PATHS.copy()

def tabs(cnt):
	cnt-=1
	c=0
	result = ''
	while not c == cnt:
		c+=1
		result += '    '
	return result

def buildChildren(ID):
	global index
	if ID in spent: return None
	spent[ID]=1
	if ID in index['folders']['ids']:
		rec = index['folders']['ids'][ID]
		if not rec['parent_id']: level = 0
		# print( index['folders']['ids'] )
		# print( ID, p, rec['title'] )
		paths = fullpath(ID)
		level = len(paths)
		if 0 and ID == '88903608dadf4da2991c1adb9ed3a8e4':
			print(level,paths)
			# paths.fullpath()
			for i in paths:
				print( tabs(len(fullpath(i))), index['folders']['ids'][i]['title'] )

		# print( ID, level, rec['title'], rec['parent_id'] )
		if False:
			try:
				print(tabs(level),rec['title'].strip(),          ' - ',index['folders']['ids'][rec['parent_id']]['title'].strip(), ID  )
			except Exception as e:
				print(tabs(level),rec['title'].strip()  )
		if ID in index['folders']['parents']:
			for iD in index['folders']['parents']:
				if not iD == ID:
					buildChildren(iD)

hierarchy = {}
_.v.IDs={}
def traverse(dic,r=0):
	if not r: _hierarchy(_.v.path)
	global index
	for k in dic:
		l=len(fullpath(k))
		# if k in _.v.IDs or not _.switches.isActive('Plus'):
		if True:
			if k in index['folders']['ids']: _.pr( l, tabs(l), index['folders']['ids'][k]['title'].strip(), c='cyan' )
			# if k in index['notes']['ids']: _.pr( l, tabs(l), index['notes']['ids'][k]['title'].strip(), c='cyan' )

			# xref = set(list(_.v.IDs.keys())) & set(list(dic[k].keys()))
			# for xr in xref:
			#     if xr in index['notes']['ids']: _.pr( l+1, tabs(l+1), index['notes']['ids'][xr]['title'].strip(), c='darkcyan' )
			#     if xr in index['folders']['ids']: _.pr( l+1, tabs(l+1), index['folders']['ids'][xr]['title'].strip(), c='darkcyan' )

			if _.switches.isActive('Notes'):
				if k in index['notes']['parents']:
					if k in index['notes']['parents']:
						xref = set(list(_.v.IDs.keys())) & set(index['notes']['parents'][k])
						for xr in xref: _.pr( l+1, tabs(l+1), index['notes']['ids'][xr]['title'].strip(), c='darkcyan' )
					# if k in index['folders']['parents']:
					#     xref = set(list(_.v.IDs.keys())) & set(index['folders']['parents'][k])
					#     for xr in xref: _.pr( l+1, tabs(l+1), index['folders']['ids'][xr]['title'].strip(), c='darkcyan' )

					# # print("index['notes']['parents'][k]",index['notes']['parents'][k])
					# for c in index['notes']['parents'][k]:
					#     # if _.showLine(index['notes']['ids'][c]['title']):
					#     if c in _.v.IDs or not _.switches.isActive('Plus'):
					#         _.pr( l+1, tabs(l+1), index['notes']['ids'][c]['title'].strip(), c='darkcyan' )
		traverse(dic[k],r=1)
_.v.print=False
def traverseSearch(dic,r=0):
	global index
	global hierarchy
	if not r: hierarchy={};_.v.IDs={};
	_.v.print=False
	for k in dic:
		if _.showLine(index['folders']['ids'][k]['title']):
			l=len(fullpath(k))
			# _.pr( l, tabs(l), index['folders']['ids'][k]['title'].strip(), c='cyan' )
		if _.switches.isActive('Notes'):
			if k in index['notes']['parents']:
				for c in index['notes']['parents'][k]:
					# if c in index['notes']['ids']:
					if _.showLine(index['notes']['ids'][c]['title']):
						_.v.IDs[c]=1
						fullpath(c)
						# _.pr( l+1, tabs(l+1), index['notes']['ids'][c]['title'].strip(), c='darkcyan' )
		traverseSearch(dic[k],r=1)

def process():
	global hierarchy
	if _.switches.isActive('Plus') or _.switches.isActive('Minus'):
		_.v.IDs={}
		if not _.switches.isActive('Search-Notes-Content'):
			for i in index['folders']['ids']: fullpath(i)
			# _.pv(hierarchy)
			traverseSearch(hierarchy.copy())
		else:
			_.pr(line=1,c='ColorBold.gray')
			queryDB(sql = "SELECT DISTINCT * FROM notes QUERY ")
			_.pr(line=1,c='ColorBold.gray')
			queryDB(sql = "SELECT DISTINCT * FROM folders QUERY ")
	_.pr(line=1,c='ColorBold.gray')
	traverse(hierarchy)
	# print(_.toYML(hierarchy))

def action():
	_.pr(line=1,c='green')
	_.pr(line=1,c='green')
	if _.switches.isActive('Search-Notes-Content'): _.switches.fieldSet( 'Notes', 'active', True )
	# if _.switches.isActive('Search-Notes-Content'): queryDB(); return None;
	load()
	all_data('folders')
	all_data('notes')

	global index

	# for k in index: print(k)
	if not _.switches.isActive('Plus') and not _.switches.isActive('Minus'):
		for i in index['folders']['ids']: fullpath(i)

	# for ID in index['top']: buildChildren(ID)
	process()
	# _.pv(index)
	# _.pr( _.toYML(_.v.IDs), c='green' )

	# _.pv(data)

def load():
	global config
	global index
	if not _.switches.isActive('Download'): index = _.getYML('pyJoplin.yml')
	token = _keychain.imp.key('joplin-token')
	# config['notes'].url = 'http://localhost:41184/notes/?token=[joplin-token]'.replace('[joplin-token]',token)

	config['folders'] = _.dot()
	config['folders'].url = 'http://localhost:41184/folders/?token=[joplin-token]'.replace('[joplin-token]',token)
	config['folders'].current = 0
	config['folders'].max = 2
	config['folders'].max = 13
	config['folders'].direction = 'ASC'
	config['folders'].direction = 'DESC'
	config['folders'].url = config['folders'].url+'&order_by=updated_time&order_dir='+config['folders'].direction+'&limit=100&page=[page]'
	config['notes'] = _.dot()
	config['notes'].url = 'http://localhost:41184/notes/?token=[joplin-token]'.replace('[joplin-token]',token)
	config['notes'].current = 0
	config['notes'].max = 2
	config['notes'].max = 100000
	config['notes'].direction = 'ASC'
	config['notes'].direction = 'DESC'
	config['notes'].url = config['notes'].url+'&order_by=updated_time&order_dir='+config['notes'].direction+'&limit=100&page=[page]'
index = {}
index['top'] = []
config = {}
data = _.dot()
_keychain = _.regImp( __.appReg, 'keychain' )
simplejson = __.imp('simplejson')
_.v.db='C:\\Users\\Scott\\.config\\joplin-desktop\\database.sqlite'
##################################################
#b)--> examples
# banner=_.Banner(dependencies)
# goss=banner.goss
# goss('-\t this app will sherlock tf out of any python app or python module')
#e)--> examples
def fields(t='notes'):
	connection = sqlite3.connect(_.v.db)
	connection.row_factory = sqlite3.Row
	sql = 'SELECT * FROM notes LIMIT 1'.replace('notes',t)
	cursor = connection.execute(sql)
	row = cursor.fetchone()
	names = row.keys()
	return names

import sqlite3
def queryDB(sql = "SELECT DISTINCT * FROM notes QUERY LIMIT 2"):
	while '  ' in sql: sql=sql.replace('  ',' ')
	table=sql.lower().split(' from ')[1].split(' ')[0]

	# sql = 'SELECT id, parend_id FROM notes LIMIT 2'

	y = "(body LIKE '%FIELD%')"
	n = "(body NOT LIKE '%FIELD%')"
	yL=[]
	yL2=[]
	nL=[]
	if table == 'notes':
		for f in _.switches.values('Plus'): yL.append(y.replace('FIELD',f))
		for f in _.switches.values('Minus'): nL.append(y.replace('FIELD',f))
	for f in _.switches.values('Plus'): yL2.append(y.replace('body','title').replace('FIELD',f))
	for f in _.switches.values('Minus'): nL.append(y.replace('body','title').replace('FIELD',f))

	add=[]
	if yL: add.append(' AND '.join(yL))
	if yL2: add.append(' AND '.join(yL2))
	if nL: add.append(' OR '.join(nL))
	if add:
		BOTH = ' ( '+' ) OR ( '.join(add)+ ' ) '
		# print(BOTH)
	else: BOTH=''
	if BOTH: BOTH = ' WHERE '+BOTH
	else: sql=sql.replace('DISTINCT','')
	sql = sql.replace('QUERY',BOTH)
	# print(sql)
	nm = fields(table)
	# ' AND '.join()
	conn = sqlite3.connect(_.v.db)
	c = conn.cursor()
	c.execute(sql)
	records = c.fetchall()
	_.pr(len(records), sql,c='purple')
	for rec in records:
		r={}
		for i,k in enumerate(nm): r[k]=rec[i]
		# _.pr( fullpath(r['parent_id']) )
		fullpath(r['id'])
		fullpath(r['parent_id'])
		_.v.IDs[r['id']]=1
		_.v.IDs[r['parent_id']]=1
		if r['parent_id']: fullpath(r['parent_id'])
		# _.pv(r)
		if _.switches.isActive('Print'):

			_.pr('    ID:',r['id'],c='yellow')
			if fullpath2(r['id']):
				_.pr(' path:',fullpath2(r['id']),c='darkcyan')
			_.pr(' path:',fullpath2(r['parent_id']),c='Background.yellow')
			_.pr(' title:',r['title'],c='Background.blue')
		if not 'body' in r:
			if _.switches.isActive('Print'): _.pr('  type: folder',c='green')
			_.v.IDs[r['id']]=1

		else:
		# if 'body' in r:
			_.v.IDs[r['id']]=1
			if _.showLine(r['body']):
				fullpath(r['parent_id'])
				if _.switches.isActive('Print'):
					body=r['body']
					for item in _.switches.values('Plus'):
						for vVv in _.caseUnspecific(r['body'],item): body=body.replace(vVv,_.pr(vVv,c='green',p=0))
					_.pr('  BODY:',body)

########################################################################################
if __name__ == '__main__':
	#b)--> examples

	# banner.pr()
	# if len(_.switches.all())==0: banner.gossip()

	#e)--> examples
	action()
	_.isExit(__file__)

# buildChildren

# process
# _.v.IDs[
# _.v.IDs=
# len(fullpath

# todo: turn stuff on
# todo: turn stuff on
# todo: turn stuff on
# todo: turn stuff on
# todo: turn stuff on
# todo: turn stuff on


import gc
import os
import cv2
import torch
import torch.nn.functional as F
import numpy as np
from tqdm import tqdm
from transformers import DistilBertTokenizer
import matplotlib.pyplot as plt
import argparse

import config as CFG
from dataset import get_transforms
from pretrain import build_loaders, make_train_valid_dfs
from CLIP import CLIPModel

def get_image_embeddings(valid_df, model_path):
    tokenizer = DistilBertTokenizer.from_pretrained(CFG.text_tokenizer)
    valid_loader = build_loaders(valid_df, tokenizer, mode="valid")

    model = CLIPModel().to(CFG.device)
    model.load_state_dict(torch.load(CFG.model_path, map_location=CFG.device))
    model.eval()

    valid_image_embeddings = []
    with torch.no_grad():
        for batch in tqdm(valid_loader):
            image_features = model.image_encoder(batch["image"].to(CFG.device))
            image_embeddings = model.image_projection(image_features)
            valid_image_embeddings.append(image_embeddings)
    return model, torch.cat(valid_image_embeddings)

def get_embeddings(text_file, image_file):
    model = CLIPModel().to(CFG.device)
    model.load_state_dict(torch.load(CFG.model_path, map_location=CFG.device))
    model.eval()

    # get text query
    f = open(text_file)
    query = f.readlines()
    query = list(map(lambda x:x.replace("\n", ""), query))

    tokenizer = DistilBertTokenizer.from_pretrained(CFG.text_tokenizer)
    encoded_query = tokenizer(query, padding=True, truncation=True, max_length=CFG.max_length)
    batch = {
        key: torch.tensor(values).to(CFG.device)
        for key, values in encoded_query.items()
    }

    # get image
    image = cv2.imread(image_file)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    transforms = get_transforms()
    image = transforms(image=image)['image']
    image = torch.tensor(image).permute(2, 0, 1).float()
    image = torch.reshape(image, ((1,) + image.shape))

    # get text & image embedding
    with torch.no_grad():
        image_features = model.image_encoder(image.to(CFG.device))
        image_embeddings = model.image_projection(image_features)
        text_features = model.text_encoder(
            input_ids=batch["input_ids"], attention_mask=batch["attention_mask"]
        )
        text_embeddings = model.text_projection(text_features)

    # find match
    image_embeddings_n = F.normalize(image_embeddings, p=2, dim=-1)
    text_embeddings_n = F.normalize(text_embeddings, p=2, dim=-1)
    dot_similarity = text_embeddings_n @ image_embeddings_n.T
    dot_similarity = dot_similarity.T
    _, indices = torch.topk(dot_similarity.squeeze(0), 3)
    matches = [query[idx] for idx in indices]
    print(matches)

def find_matches(text_embeddings, image_embeddings, query):
    image_embeddings_n = F.normalize(image_embeddings, p=2, dim=-1)
    text_embeddings_n = F.normalize(text_embeddings, p=2, dim=-1)
    dot_similarity = text_embeddings_n @ image_embeddings_n.T

    _, indices = torch.topk(dot_similarity.squeeze(0), 3)
    matches = [image_filenames[idx] for idx in indices[::5]]

    _, axes = plt.subplots(3, 3, figsize=(10, 10))
    for match, ax in zip(matches, axes.flatten()):
        image = cv2.imread(f"{CFG.image_path}/{match}")
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        ax.imshow(image)
        ax.axis("off")

    plt.show()
    query_ = query.replace(" ","_")
    result_filename = query_ + "_" + CFG.model_name + "_" + CFG.dataset + "_result.png"
    plt.savefig(result_filename)

def main(args):
    if (args.query_text is None) or (args.query_image is None):
        raise Exception("Usage :[python3 inference_image_to_text.py --query_text={YOUR QUERY TEXT FILE} --query_image={YOUR QUERY IMAGE FILE}]")

    get_embeddings(args.query_text, args.query_image)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--query_text', type=str)
    parser.add_argument('--query_image', type=str)
    args = parser.parse_args()
    main(args)


#the method follows Eqn. 7 in Zhou's paper
#the variance of the chi-square minimization fitting of the antisymmetric cross power spectrum dipole

import numpy as np
from scipy import interpolate, integrate
from scipy.stats import chi2
from scipy.special import erfinv, erfcinv
#*****************************#
## Cosmological Parameters ##
#*****************************#
HubblePara = 0.678
## no unit
HubbleLength = 3000
## in unit Mpc/h

OMega_m = 0.306
OMega_r = 8.6e-5
OMega_lambda = 1.0-OMega_m-OMega_r
## Constants of LambdaCDM, no units

nu_21 = 1420.405752
nu_CO = 115270
## Frequency (MHz) of emission lines at rest

lambda_21 = 0.2110611405
lambda_CO = 0.00261
## Wavelength (meters) of emission lines at rest

def E(z):
    """
    The scale factor which relates the current Hubble parameter H0 with the Hubble parameter at redshift z H(z) by H(z) = E(z) H0.
    """
    return np.sqrt( OMega_m*(1+z)**3 + OMega_r*(1+z)**4 + OMega_lambda )

def X(z):
    """
    The radial co-moving distance (Mpc/h) at redshift z.
    """
    integral = integrate.quad(lambda x: 1./E(x), 0, z)[0]
    integral *= HubbleLength
    return integral

def Y(z, nu_rest):
    """
    Y = d r_parallel (Mpc/h) / d v (MHz) = Hubblelength * (1+z)^2/E(z)/nu_rest.

    """
    return HubbleLength * (1.+z) **2. / E(z) / nu_rest

#*****************************#
## Main Functions ##
#*****************************#

def significance_level(signal_noise_ratio_square_sum, free_degree):
    """
    Calculate the significance level for claiming a detection.
    Assume the square sum of signal-noise-ratio from multiple measurements follows a chi-square distribution,
    then get the p-value and convert it to the equivalent significance level of a normal distribution.
    """
    p = chi2.cdf(signal_noise_ratio_square_sum, free_degree)
    significance = erfinv(p)
    if significance == float('inf'):
        p_comp = integrate.quad(lambda x: chi2.pdf(x, free_degree), signal_noise_ratio_square_sum, np.inf, epsrel = 1e-5)[0]
        significance = erfcinv(p_comp)
    return significance

def reliable_k_range(z, Omega_beam_CO, Omega_patch, delta_nu_CO, delta_D, L_min_21, L_max_21, delta_nu_21):
    """
    Omega_beam_CO : [radian]^2, the instant solid-angle on the sky subtented by one beam of CO detection
    Omega_patch : [radian]^2, the solid-angle on the sky subtented by a box
    delta_nu_CO : MHz, the width of single frequency channel of CO detection
    delta_D : Mpc/h, the length of the observation box along the LoS direction
    L_min_21: meters, the minimum baseline length of 21cm observation
    L_max_21: meters, the maximum baseline length of 21cm observation
    delta_nu_21 : MHz, the width of single frequency channel of 21 detection
    """
    k_perp_max_CO = 2.*np.pi / (X(z) * np.sqrt(Omega_beam_CO))
    k_perp_min_CO = 2.*np.pi / (X(z) * np.sqrt(Omega_patch))
    k_para_max_CO = 2.*np.pi / (Y(z, nu_CO) * delta_nu_CO)
    k_perp_max_21 = 2. * np.pi * L_max_21 / lambda_21 / (1 + z) / X(z);
    k_perp_min_21 = 2. * np.pi * L_min_21 / lambda_21 / (1 + z) / X(z);
    k_para_max_21 = 2. * np.pi / (Y(z, nu_21) * delta_nu_21)

    k_perp_max = min(k_perp_max_CO, k_perp_max_21);     k_perp_min = max(k_perp_min_CO, k_perp_min_21)
    k_para_max = min(k_para_max_CO, k_para_max_21);     k_para_min = 2.*np.pi / delta_D
    return [k_perp_min, k_perp_max, k_para_min, k_para_max] # in the unit of h/Mpc

def W(k_perp, k_para, z, nu_rest, Omega_beam, delta_nu):
    """
    The modified factor due to limited resolution
    -----------
    Parameters:
    -----------
    k_perp : h/Mpc, k perpendicular to the LoS direction
    k_para : h/Mpc, k parallel to the LoS direction
    z : redshift
    nu_rest : MHz, the frequency of the emission line at its rest frame
    Omega_beam : [radian]^2, the instant solid-angle on the sky subtented by one beam
    delta_nu : MHz, the width of single frequency channel
    """
    sigma_perp = X(z)*np.sqrt(Omega_beam) / (2 * np.pi)
    sigma_para = Y(z, nu_rest)*delta_nu / (2 * np.pi)
    modification = np.exp(sigma_perp**2 * k_perp**2 + sigma_para**2 * k_para**2)
    return modification

def error_CO_auto(k, delta_k, theta, delta_theta, P_k, z, Tsys, Omega_patch, NUM_PATCH, delta_D, Omega_beam, delta_nu, Nfeeds, t_int):
    """
    Calculate the cosmic variance on CO auto power spectrum at certain (k, theta)
    -----------
    Parameters:
    -----------
    k : h/Mpc
    delta_k : h/Mpc
    theta: radians
    delta_theta: radians
    P_k : [muK]^2[Mpc/h]^3, the auto power spectrum of CO line
    z : redshift
    Tsys : muK
    Omega_patch : [radian]^2, the solid-angle on the sky subtented by a box
    NUM_PATCH : number of patches in a survey
    delta_D : Mpc/h, the length of the observation box along the LoS direction
    Omega_beam : [radian]^2, the instant solid-angle on the sky subtented by one beam
    delta_nu : MHz, the width of single frequency channel
    Nfeeds : the number of independent feeds, like different beams, polarizations ...
    t_int : second, the total integration time
    --------
    Returns:
    --------
    sigma_error: [muK]^2[Mpc/h]^3
    sigma_cosmic_variance : [muK]^2[Mpc/h]^3
    sigma_noise : [muK]^2[Mpc/h]^3
    Nmodes :
    """

    Omega_survey = NUM_PATCH * Omega_patch
    V_survey = X(z)**2*Omega_survey*delta_D
    t_int *= 1e6

    k_perp = k * np.sin(theta); k_para = k * np.cos(theta)
    N_modes = V_survey / (4. * np.pi**2) * k**2 * delta_k * np.sin(theta) * delta_theta
    sigma_noise = W(k_perp, k_para, z, nu_CO, Omega_beam, delta_nu) * X(z)**2*Y(z, nu_CO)*Tsys**2*Omega_survey/Nfeeds/t_int/N_modes**0.5
    sigma_cosmic_variance = P_k / N_modes**0.5
    sigma_error = sigma_cosmic_variance + sigma_noise

    return sigma_error, sigma_cosmic_variance, sigma_noise, N_modes

def n_u(u, wavelength):
    """
    The number density of redundant baselines in uv space, assuming a SKA-low core configuration.
    u : no unit, L/wavelength
    wavelength : m
    """
    baseline_number_density = [0.0609823,0.05809261,0.05520293,0.05231325,0.04942356,\
    0.04673317,0.04404277,0.04135238,0.03866198,0.03597158,0.03328119,0.03079008,\
    0.02829897,0.02580787,0.02331676,0.02102494,0.01873312,0.01644131,0.01414949,\
    0.01205696,0.00996443,0.00807119,0.00617795,0.00468328,0.00318862,0.00209253,0.00099644,0.00049822,0.]
    baseline_length = [ 35, 70, 105,  140,  175,  210,  245,  280,  315,  350,  385, 420,  455,  \
    490,  525,  560,  595,  630,  665,  700,  735,  770, 805,  840,  875,  910,  945,  980, 1015 ]
    f = interpolate.interp1d(baseline_length, baseline_number_density, kind = "cubic", fill_value="extrapolate")
    return f(u*wavelength)*wavelength**2

def error_21_auto(k, delta_k, theta, delta_theta, P_k, z, Tsys, Omega_patch, NUM_PATCH, Omega_beam, delta_D, Nfeeds, t_int):
    """
    Calculate the error bar on 21cm auto power spectrum at certain (k, theta)
    -----------
    Parameters:
    -----------
    k : h/Mpc
    delta_k : h/Mpc, the width of the k-shell
    theta: radians
    delta_theta: radians
    P_k : [muK]^2[Mpc/h]^3
    z : redshift
    Tsys : muK
    delta_k : h/Mpc, the width of the k-shell
    Omega_patch : [radian]^2, the solid-angle on the sky subtented by a box
    NUM_PATCH : number of patches in a survey
    Omega_beam : [radian]^2, the instant solid-angle on the sky subtented by one beam
    delta_D : Mpc/h, the length of the observation box along the LoS direction
    Nfeeds : the number of independent feeds, like different beams, polarizations ...
    t_int : second, the total integration time
    --------
    Returns:
    --------
    sigma_error: [muK]^2[Mpc/h]^3
    sigma_cosmic_variance : [muK]^2[Mpc/h]^3
    sigma_noise : [muK]^2[Mpc/h]^3
    Nmodes :
    """
    Omega_survey = NUM_PATCH * Omega_patch
    V_survey = X(z)**2*Omega_survey*delta_D
    t_int *= 1e6

    k_perp = k * np.sin(theta); k_para = k * np.cos(theta)
    u_perp = k_perp*X(z)/2/np.pi
    N_modes = V_survey / (4. * np.pi**2) * k**2 * delta_k * np.sin(theta) * delta_theta
    sigma_noise = X(z)**2*Y(z, nu_21)*Tsys**2*Omega_patch*NUM_PATCH*Omega_beam/t_int/Nfeeds/n_u(u_perp, lambda_21*(1+z)) / N_modes**0.5
    sigma_cosmic_variance = P_k / N_modes**0.5
    sigma_error = sigma_cosmic_variance + sigma_noise

    return sigma_error, sigma_cosmic_variance, sigma_noise, N_modes

def error_21CO_SymCross(k, delta_k, theta, delta_theta, PS_k, z, Omega_patch, NUM_PATCH, delta_D):
    """
    -----------
    Parameters:
    -----------
    k : h/Mpc
    delta_k : h/Mpc, the width of the k-shell
    theta: radians
    delta_theta: radians
    PS_k : [muK]^2[Mpc/h]^3, the symmetric cross power spectrum between HI and CO lines
    z : redshift
    Omega_patch : [radian]^2, the solid-angle on the sky subtented by a box
    NUM_PATCH : number of patches in a survey
    delta_D : Mpc/h, the length of the observation box along the LoS direction
    --------
    Returns:
    --------
    sigma_error: [muK]^2[Mpc/h]^3
    Nmodes :
    """
    V_survey = X(z)**2 * Omega_patch * NUM_PATCH * delta_D

    N_modes = V_survey / (4. * np.pi**2) * k**2 * delta_k * np.sin(theta) * delta_theta
    sigma_error = PS_k / N_modes**0.5

    return abs(sigma_error), N_modes

def error_21CO_AntiCross(k, delta_k, theta, delta_theta, PA_k, z, Omega_patch, NUM_PATCH, delta_D):
    """
    -----------
    Parameters:
    -----------
    k : h/Mpc
    delta_k : h/Mpc, the width of the k-shell
    theta: radians
    delta_theta: radians
    PA_k : [muK]^2[Mpc/h]^3, the anti-symmetric cross power spectrum between HI and CO lines
    z : redshift
    Omega_patch : [radian]^2, the solid-angle on the sky subtented by a box
    NUM_PATCH : number of patches in a survey
    delta_D : Mpc/h, the length of the observation box along the LoS direction
    --------
    Returns:
    --------
    sigma_error: [muK]^2[Mpc/h]^3
    Nmodes :
    """
    V_survey = X(z)**2 * Omega_patch * NUM_PATCH * delta_D
    PI_k = PA_k * (3 / 4 / np.pi)**0.5 * np.cos(theta)

    N_modes = V_survey / (4. * np.pi**2) * k**2 * delta_k * np.sin(theta) * delta_theta
    sigma_error = PI_k / N_modes**0.5

    return sigma_error, N_modes

def VAR_PA_k(k, delta_k, z, Omega_patch, NUM_PATCH, Tsys_CO, Tsys_21, Omega_beam_CO, Omega_beam_21, delta_nu_CO, delta_nu_21, Nfeeds_CO, Nfeeds_21, t_int_CO, t_int_21, delta_D, L_min_21, L_max_21, P_k_CO, P_k_21, PS_k, PA_k):
    """
    Calculate the variance on the dipole of the CO auto antisymmetric power spectrum at certain k
    -----------
    Parameters:
    -----------
    k : h/Mpc
    delta_k : h/Mpc, the width of the k-shell
    z : the redshift
    z : redshift

    Omega_patch : [radian]^2, the solid-angle on the sky subtented by a box
    NUM_PATCH : number of patches in a survey
    Tsys_CO, Tsys_21 : muK, the system temperature of the two lines detction respectively
    Omega_beam_CO, Omega_beam_21 : radians^2, the solid angle of CO and 21cm detection
    delta_nu_CO, delta_nu_21 : MHz, the width of single frequency channel of CO and 21cm lines
    N_feeds_CO, N_feeds_21 : the number of independent feeds, like different beams, polarizations ...
    t_int_CO, t_int_21 : second, the total integration time
    delta_D : Mpc/h, the length of the observation box along the LoS direction
    L_min_21, L_max_21: meters, the minimum and the maximum baseline length of 21cm observation
    P_k_CO, P_k_21, PS_k, PA_k : [muK]^2[Mpc/h]^3, the auto power spectrum of CO line, the auto power spectrum of 21cm line, the symmetric cross power spectrum and the antisymmetric power spectru,
    --------
    Returns:
    --------
    sigma_error: [muK]^2[Mpc/h]^3, the total error the dipole detection
    sigma_cosmic_variance : [muK]^2[Mpc/h]^3, the cosmic variance of the dipole detection
    """

    #the limits are around 0.033, 0.56, 0.024, 143 h/Mpc
    k_perp_min, k_perp_max, k_para_min, k_para_max = reliable_k_range(z, Omega_beam_CO, Omega_patch, delta_nu_CO, delta_D, L_min_21, L_max_21, delta_nu_21)
    if k**2 < k_perp_min**2 + k_para_min**2 or k**2 > k_perp_max**2 + k_para_max**2:
        return 'ERROR: k is out of the reliable range'
    delta_theta = k_para_min / k
    sin_max = k_perp_max / k if k_perp_max / k < 1 else 0.999
    theta_max = np.arcsin(sin_max)
    theta, VAR_PA_k, VAR_PA_k_cv = 0.01, 0, 0
    while theta < theta_max:
        #the total error and the cosmic variance
        sigma_P_k_CO, sigma_P_k_CO_cv, _, _ = error_CO_auto(k, delta_k, theta, delta_theta, P_k_CO, z, Tsys_CO, Omega_patch, NUM_PATCH, delta_D, Omega_beam_CO, delta_nu_CO, Nfeeds_CO, t_int_CO)
        sigma_P_k_21, sigma_P_k_21_cv, _, _ = error_21_auto(k, delta_k, theta, delta_theta, P_k_21, z, Tsys_21, Omega_patch, NUM_PATCH, Omega_beam_21, delta_D, Nfeeds_21, t_int_21)
        sigma_PS_k, _ = error_21CO_SymCross(k, delta_k, theta, delta_theta, PS_k, z, Omega_patch, NUM_PATCH, delta_D)
        sigma_PI_k, _ = error_21CO_AntiCross(k, delta_k, theta, delta_theta, PA_k, z, Omega_patch, NUM_PATCH, delta_D)
        #print(sigma_P_k_CO, sigma_P_k_21, sigma_PS_k, sigma_PI_k)
        VAR_PI_k_theta_cv = sigma_P_k_CO_cv * sigma_P_k_21_cv - sigma_PS_k**2 + sigma_PI_k**2
        VAR_PI_k_theta = sigma_P_k_CO * sigma_P_k_21 - sigma_PS_k**2 + sigma_PI_k**2
        VAR_PA_k_cv += ((3 / 4 / np.pi * np.cos(theta)**2) / VAR_PI_k_theta_cv)
        VAR_PA_k += ((3 / 4 / np.pi * np.cos(theta)**2) / VAR_PI_k_theta) #eqn 8 in Zhou et al. 2021
        theta += delta_theta
    VAR_PA_k = 1 / VAR_PA_k; VAR_PA_k_cv = 1 / VAR_PA_k_cv
    return VAR_PA_k**0.5, VAR_PA_k_cv**0.5 #return the sigma of error


from collections import defaultdict
class Solution:
    # O(n)
    def longestConsecutive(self, nums: List[int]) -> int:
        result = 0
        d = self.get_hashmap(nums)
        for num in nums:
            if (num-1) in d:
                continue
            else:
                count = 1 # One Because this number already counted
                number = num # temp variable
                # Run loop till next number in sequence exists
                while (number+1) in d:
                    number += 1
                    count += 1
                result = max(result,count)
        return result

    def get_hashmap(self, nums):
        d = defaultdict(int)
        for num in nums:
            d[num] = d[num] + 1
        return d

    # O(nlogn)
    def longestConsecutiveOld(self, nums: List[int]) -> int:
        if len(nums) == 0:
            return 0
        nums.sort()
        index = 0
        maxLength = -1
        count = 1
        # print('Sorted', nums)
        while index < len(nums)-1:
            if nums[index+1] == nums[index]:
                index += 1
            elif nums[index+1] == nums[index]+1:
                count += 1
                index += 1
            else:
                maxLength = max(maxLength, count)
                count = 1
                index += 1
        maxLength = max(maxLength, count)
        return maxLength


# -*- coding: utf-8 -*-
"""
Setup file for EDWIN
"""
import os
from setuptools import setup, find_packages
import pkg_resources

repo = os.path.dirname(__file__)
try:
    from git_utils import write_vers
    version = write_vers(vers_file='ed_win/__init__.py', repo=repo, skip_chars=1)
except Exception:
    version = '999'

try:
    from pypandoc import convert_file

    def read_md(f): return convert_file(f, 'rst', format='md')
except ImportError:
    print("warning: pypandoc module not found, could not convert Markdown to RST")

    def read_md(f): return open(f, 'r').read()

setup(name='ed_win',
      version=version,
      description='EDWIN an optimization and design package for electrical networks in windfarms',
      long_description=read_md('README.md'),
      url='https://gitlab.windenergy.dtu.dk/TOPFARM/EDWIN',
      author='DTU Wind Energy',
      author_email='juru@dtu.dk',
      license='MIT',
      packages=find_packages(),
      install_requires=[
          'matplotlib',  # for plotting
          'numpy',  # for numerical calculations
          'xarray',  # for WaspGridSite data storage
          'scipy',  # constraints
      ],
      extras_require={
          'test': [
              'pytest',  # for testing
              'pytest-cov',  # for calculating coverage
              'sphinx',  # generating documentation
              'sphinx_rtd_theme',  # docs theme
          ]},
      zip_safe=True)


#!/usr/bin/env python
# coding: utf-8

# In[156]:

import streamlit as st
import pickle
import pandas as pd

# In[162]:

#links = pd.read_csv("/Users/hi/Desktop/data_science_wbs_school/Recommender_system/project/c1ba2d8cbaa22297e5d9b0b7a17fcb7awbsflix-dataset/ml-latest-small/links.csv")

#movies = pd.read_csv("/Users/hi/Desktop/data_science_wbs_school/Recommender_system/project/c1ba2d8cbaa22297e5d9b0b7a17fcb7awbsflix-dataset/ml-latest-small/movies.csv")
#ratings = pd.read_csv("/Users/hi/Desktop/data_science_wbs_school/Recommender_system/project/c1ba2d8cbaa22297e5d9b0b7a17fcb7awbsflix-dataset/ml-latest-small/ratings.csv")
#tags = pd.read_csv("/Users/hi/Desktop/data_science_wbs_school/Recommender_system/project/c1ba2d8cbaa22297e5d9b0b7a17fcb7awbsflix-dataset/ml-latest-small/tags.csv")

# In[163]:

movies_dict = pickle.load(open('movies_dict.pkl', 'rb'))
movies= pd.DataFrame(movies_dict)

rating_final = pickle.load(open('rating_final.pkl', 'rb'))
ratings = pickle.load(open('ratings.pkl', 'rb'))
rating_avg = pickle.load(open('rating_avg.pkl', 'rb'))
rating_count = pickle.load(open('rating_count.pkl', 'rb'))
users = pickle.load(open('users.pkl', 'rb'))
user_Id = pickle.load(open('user_Id.pkl', 'rb'))
movie_rating_tab = pickle.load(open('movie_rating_tab.pkl', 'rb'))
rating_final1 = pickle.load(open('rating_final1.pkl', 'rb'))

# In[180]:

st.set_page_config(page_title="WEBFLIX", layout="wide")
st.title('WBSFLIX')

#search_choices = ['Movie title', 'Genre', 'Rating', 'Tags']
#search_selected = st.sidebar.selectbox("Your choice please", search_choices)
movie_name = movies['title'].values

#dropdown_movie = st.selectbox('What would you like to watch today?', movie_name)

selected_movie_name= st.selectbox('What would you like to watch today?',movie_name)

st.write('You selected:', selected_movie_name)

def get_recommendations (movie_name,n=20):
    movie_rating_tab= pd.pivot_table(data=rating_final, values='rating', index='userId', columns='title')
      #movie_rating_tab = movie_rating_tab.fillna(0)
    rest_movie_ratings = movie_rating_tab[movie_name]
    similar_to_rest = movie_rating_tab.corrwith(rest_movie_ratings)
    corr_rest = pd.DataFrame(similar_to_rest, columns = ['Pearson_coeif'])
    corr_rest.dropna(inplace=True)
    corr_rest.drop(movie_name, inplace=True)
    similar = corr_rest.merge(rating_final, on = "title")
    similar = similar.drop(['rating', 'timestamp', 'userId'], axis=1)
    similar = similar.drop_duplicates("title")
    similar_final= similar[similar['rating_count']>=100].sort_values(by = 'Pearson_coeif', ascending = False).head(n)
    similar_final = similar_final.drop(["overall_rating","rating_count","Pearson_coeif"], axis=1)
    similar_final["rating_avg"] = round(similar_final["rating_avg"], 1)
    similar_final["Year"]= similar_final["title"].str.findall("\d{4}")
    similar_final["title"]= similar_final["title"].str.replace('(\(\d{4})\)', '')
    similar_final = similar_final[['movieId', 'title', 'Year', 'genres', 'rating_avg']]
    similar_final_new= similar_final.set_index("movieId")
    return(similar_final_new)

with st.container():
    st.subheader('Top picks for you based on the movie you selected :')
    buffer1, col1, buffer2 = st.columns([2.5, 3, 1])
    is_clicked = col1.button(label = 'Recommended')
    if is_clicked :
        recommendations = get_recommendations(selected_movie_name)
        columns = st.columns((1, 4, 1))
        columns[1].write(recommendations)

def popular_table (movie_count):

  rating_avg = pd.DataFrame(ratings.groupby('movieId')['rating'].mean())
  rating_avg['overall_rating'] = round((ratings.groupby('movieId')['rating'].mean()))
  rating_count = pd.DataFrame(ratings.groupby('movieId')['rating'].count())
  rating_final  = rating_avg.merge(rating_count, on = 'movieId').merge(movies, on = "movieId")
  rating_final.rename(columns = {'rating_x':'rating_avg', 'rating_y':'rating_count'}, inplace = True)
  rating_final = rating_final[rating_final['rating_count']>=100].sort_values("rating_count", ascending=False)
  rating_final = rating_final.drop_duplicates(["title"])
  top_movies = rating_final.head(movie_count)

  return(top_movies)

# list of top 15 popular movies in our database
popular_movies = pickle.load(open('popular_all.pkl', 'rb'))

#session.options = st.multiselect(label="Select Movies", options=popular_movies)

#dataframe = None

#session.slider_count = st.slider(label="movie_count", min_value=5, max_value=138)

#with st.container():
   # is_clicked = col1.button(label="Popular")
    #if is_clicked:
       # dataframe = popular_table(session.options, movie_count = session.slider_count, data= popular_movies)
   # if dataframe is not None:
        #st.table(dataframe)

with st.container():
    st.subheader('Top 20:')
    buffer1, col1, buffer2 = st.columns([2.8, 3, 1])
    is_clicked = col1.button(label = "Popular")
    if is_clicked:
        columns = st.columns((0.5, 0.8, 0.3))
        columns[1].write(popular_movies)

def special_picks_for_you(user_id, n=30):
    user_rating_tab = movie_rating_tab.T.fillna(0)
  # computing cosine similarities
    from sklearn.metrics.pairwise import cosine_similarity
    user_similarities = pd.DataFrame(cosine_similarity(user_rating_tab),
                                  columns=user_rating_tab.index,
                                  index=user_rating_tab.index)
  # computing the weights for the users
    weights = (
        user_similarities.query("userId!=@user_id")[user_id] / sum(user_similarities.query("userId!=@user_id")[user_id])
            )
  # select movies that the inputed user has not rated, # dataframe for the non rated movies by all the user
    not_rated_movies = user_rating_tab.loc[user_rating_tab.index!=user_id, user_rating_tab.loc[user_id,:]==0]

    weighted_averages = pd.DataFrame(not_rated_movies.T.dot(weights), columns=["predicted_rating"])
    recommendations = weighted_averages.merge(movies, left_index=True, right_on="title").merge(rating_final1, on ="movieId").merge(users, on ="movieId").drop_duplicates("title")
    special_picks = recommendations[recommendations['rating_count']>100].sort_values("predicted_rating", ascending=False).head(n)
    special_picks = special_picks.drop(["userId","overall_rating","rating_count"], axis=1)
    special_picks["rating_avg"] = round(special_picks["rating_avg"], 1)
    special_picks["predicted_rating"] = round(special_picks["predicted_rating"], 1)

    special_picks["Year"]= special_picks["title"].str.findall("\d{4}")
    special_picks["title"]= special_picks["title"].str.replace('(\(\d{4})\)', '')
    special_picks_new = special_picks[['movieId', 'title', 'Year', 'genres', 'rating_avg', 'predicted_rating']]

    special_picks_new = special_picks_new.set_index('movieId')
    return(special_picks_new)

with st.container():
    st.subheader('Search by user Id:')
    user_id = user_Id
    user_choices =st.selectbox('See what other user have liked by User Id', user_id)
    buffer1, col1, buffer2 = st.columns([3, 3, 1])
    is_clicked = col1.button(label = 'Show')
    if is_clicked :
        special_picks = special_picks_for_you(user_choices)
        columns = st.columns((1, 6.8, 1))
        columns[1].write(special_picks)

#import streamlit.components.v1 as components

# In[ ]:

# In[ ]:

# In[ ]:

# In[ ]:


import pytest
from src.Ej_4_3 import operaciones_conjuntos

@pytest.mark.parametrize(
    "inSet1, inSet2, outTupla",
    [  # interseccion, solo_prim, solo_seg
        ({"manzana", "pera", "naranja", "plátano", "uva"},
         {"manzana", "pera", "durazno", "sandía", "uva"},
         ({"manzana", "pera", "uva"}, {"naranja", "plátano"}, {"durazno", "sandía"})),
        ({"manzana", "kiwi", "naranja", "plátano", "uva"},
         {"manzana", "pera", "durazno", "sandía", "uva"},
         ({"manzana", "uva"}, {"kiwi", "naranja", "plátano"}, {"pera", "durazno", "sandía"})),
    ]
)
def test_operaciones_conjuntos(inSet1, inSet2, outTupla):
    assert operaciones_conjuntos(inSet1, inSet2) == outTupla


#!/usr/bin/env python3
from astropy.io import fits
import numpy as np
import matplotlib.pyplot as plt

def compare(fname_trace, fname_img=None, fname_spec=None, figname=None, dpi=240, fname_bpm=None, vmin=None, vmax=None):
    """ compare img and trace """
    trace = fits.open(fname_trace)
    if fname_img:
        img = fits.open(fname_img)
        linecol = "w"
    else:
        linecol = "k"
    if fname_spec:
        spec = fits.open(fname_spec)
    else:
        spec = None
    if fname_bpm:
        bpm = fits.open(fname_bpm)
    else:
        bpm = None

    X = np.arange(2048)
    FIG = plt.figure(figsize=(10, 4))

    for i in [1, 2, 3]:
        ax = FIG.add_subplot(1, 3, i)
        ax.set_xticks([])
        ax.set_yticks([])

        try:
            tdata = trace[i].data
        except:
            print("extension %s is missing, skipping." % i)
            continue
        if tdata is None:
            print("Data for CHIP%s is empty, skipping." % i)
            continue

        if fname_img:
            imgdata = img["CHIP%d.INT1" % i].data
            if imgdata is not None:
                # imgdata = np.ma.masked_where(np.isnan(imgdata),imgdata)
                # imgdata = np.nan_to_num(imgdata)
                if bpm is not None:
                    bpmdata = bpm["CHIP%d.INT1" % i].data
                    if bpmdata is not None:
                        imgdata[bpmdata != 0] = np.nan
                if vmin is None:
                    _vmin = np.nanpercentile(imgdata, 5)
                else:
                    _vmin = vmin
                if vmax is None:
                    _vmax = np.nanpercentile(imgdata, 98)
                else:
                    _vmax = vmax
                ax.imshow(imgdata, origin="lower", cmap="plasma", vmin=_vmin, vmax=_vmax, interpolation="none")

        for t in tdata:
            pol = np.polyval(t["Upper"][::-1], X)
            ax.plot(X, pol, ":" + linecol)
            pol = np.polyval(t["Lower"][::-1], X)
            ax.plot(X, pol, ":" + linecol)
            pol = np.polyval(t["All"][::-1], X)
            ax.plot(X, pol, "--" + linecol)
            if np.isnan(pol[1024]):
                continue
            ax.text(
                1024,
                pol[1024],
                "order: %s     trace: %s" % (t["order"], t["TraceNb"]),
                color=linecol,
                horizontalalignment="center",
                verticalalignment="center",
                size=8,
            )
        ax.axis((1, 2048, 1, 2048))

    FIG.tight_layout(pad=0.02)
    if figname is None:
        figname = fname_trace.replace(".fits", ".png")
    print(f"Saving plot: {figname}")
    plt.savefig(figname, dpi=dpi)

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("trace")
    parser.add_argument("img")
    parser.add_argument("--spec")
    parser.add_argument("--bpm")
    parser.add_argument("--out")
    parser.add_argument("--dpi", default=1000)
    parser.add_argument("--vmin", type=float)
    parser.add_argument("--vmax", type=float)
    args = parser.parse_args()

    fname_trace = args.trace
    fname_img = args.img
    fname_spec = args.spec
    fname_out = args.out
    fname_bpm = args.bpm
    dpi = args.dpi
    vmin = args.vmin
    vmax = args.vmax

    compare(fname_trace, fname_img, fname_spec, fname_out, fname_bpm=fname_bpm, dpi=dpi, vmin=vmin, vmax=vmax)


#!/usr/bin/python3
'''gets subscribers of subreddit '''
import requests

def number_of_subscribers(subreddit):
    '''does the thing'''
    url = 'https://www.reddit.com/r/{}/about.json'.format(subreddit)
    headers = {'User-Agent': 'Python/1.0(Holberton Project)'}
    request = requests.get(url, headers=headers, allow_redirects=False)
    try:
        return request.json()['data']['subscribers']
    except Exception:
        return 0


import pytz
import json
import hashlib
import pandas as pd
import quantutils.dataset.pipeline as ppl
import marketinsights.utils.http as http
from marketinsights.remote.models import MIModelServer
from marketinsights.remote.ibmcloud import CloudFunctions
from marketinsights.utils.auth import CredentialsStore

class MIAssembly():

    def __init__(self, modelSvr=None, credentials_store=None, secret=None):

        if not credentials_store:
            credentials_store = CredentialsStore()

        self.fun = CloudFunctions(credentials_store)
        if not secret:
            secret = "marketinsights-local-cred"
        self.credentials = credentials_store.getSecrets(secret)
        self.modelSvr = modelSvr

    def put_dataset(self, data, dataset_desc, market, debug=False):
        dataset = Dataset.csvtojson(data, dataset_desc, market)
        headers = {
            'X-IBM-Client-Id': self.credentials["clientId"],
            'X-IBM-Client-Secret': self.credentials["clientSecret"],
            'content-type': 'application/json'
        }
        url = "".join([self.credentials["mi-api-endpoint"], "/v1/datasets"])
        return http.put(url=url, headers=headers, data=dataset, debug=debug)

    def get_dataset(self, dataset_desc, market, debug=False):
        return self.get_dataset_by_id(Dataset.generateMarketId(dataset_desc, market), debug)

    def get_dataset_by_id(self, dataset_id, debug=False):
        headers = {
            'X-IBM-Client-Id': self.credentials["clientId"],
            'X-IBM-Client-Secret': self.credentials["clientSecret"],
            'accept': 'application/json'
        }

        # TODO eliminate id, no need.
        query = {
            'where': {
                'id': dataset_id,
            }
        }
        url = "".join([self.credentials["mi-api-endpoint"], "/v1/datasets?filter=", json.dumps(query)])
        resp = http.get(url=url, headers=headers, debug=debug)
        dataset = resp[0]
        return [Dataset.jsontocsv(dataset), dataset["dataset_desc"]]

    # TODO  (without getting data too)
    def get_dataset_desc(self):
        pass

    def put_model(self, data, debug=False):
        headers = {
            'X-IBM-Client-Id': self.credentials["clientId"],
            'X-IBM-Client-Secret': self.credentials["clientSecret"],
            'content-type': 'application/json'
        }
        url = "".join([self.credentials["mi-api-endpoint"], "/v1/models"])
        return http.put(url=url, headers=headers, data=json.dumps(data), debug=debug)

    def get_model(self, modelId, debug=False):
        headers = {
            'X-IBM-Client-Id': self.credentials["clientId"],
            'X-IBM-Client-Secret': self.credentials["clientSecret"],
            'accept': 'application/json'
        }
        url = "".join([self.credentials["mi-api-endpoint"], "/v1/models/", modelId])
        return http.get(url=url, headers=headers, debug=debug)

    def put_training_run(self, data, debug=False):
        headers = {
            'X-IBM-Client-Id': self.credentials["clientId"],
            'X-IBM-Client-Secret': self.credentials["clientSecret"],
            'content-type': 'application/json'
        }
        url = "".join([self.credentials["mi-api-endpoint"], "/v1/training_runs"])
        return http.put(url=url, headers=headers, data=json.dumps(data), debug=debug)

    def get_training_run(self, training_run_id, debug=False):
        headers = {
            'X-IBM-Client-Id': self.credentials["clientId"],
            'X-IBM-Client-Secret': self.credentials["clientSecret"],
            'accept': 'application/json'
        }
        url = "".join([self.credentials["mi-api-endpoint"], "/v1/training_runs/", training_run_id])
        return http.get(url=url, headers=headers, debug=debug)

    def get_predictions_with_dataset(self, dataset, training_run_id, labels=1, debug=False):

        if not self.modelSvr:
            raise Exception("No Model Server configured")

        features = dataset.iloc[:, :-labels]
        # Send the dataset features to the model and retrieve the scores (predictions)
        return pd.concat([dataset, self.modelSvr.getPredictions(features, training_run_id, debug=debug)], axis=1)

    def get_predictions_with_dataset_id(self, dataset_id, training_run_id, start=None, end=None, debug=False):

        # Get the dataset from storage, crop and strip out labels
        dataset, dataset_desc = self.get_dataset_by_id(dataset_id, debug)

        if debug:
            print(dataset)

        if dataset.empty:  # No predictions in this time range
            return None

        return self.get_predictions_with_dataset(dataset, training_run_id, labels=dataset_desc["labels"], debug=debug)

    def get_predictions_with_raw_data(self, data, training_id, debug=False):

        training_run = self.get_training_run(training_id)
        if debug:
            print("Training run : " + str(training_run))

        dataset_id = training_run["datasets"][0]
        _, dataset_desc = self.get_dataset_by_id(dataset_id)
        pipeline = dataset_desc["pipeline"]

        # Generate a dataset from the raw data through the given pipeline
        dataset = Dataset().createDataset(pipeline["id"], data, pipeline["pipeline_desc"], debug)

        if dataset.empty:
            return dataset

        #dataset = dataset.iloc[:, :-dataset_desc["labels"]]
        if debug:
            print("Sending feature vector : " + str(dataset))

        return self.get_predictions_with_dataset(dataset, training_id, labels=dataset_desc["labels"], debug=debug)

    @staticmethod
    def generateMarketId(dataset_desc, market):
        return hashlib.md5("".join([market, json.dumps(dataset_desc["pipeline"], sort_keys=True), str(dataset_desc["features"]), str(dataset_desc["labels"])]).encode('utf-8')).hexdigest()

    @staticmethod
    def generateTrainingId(dataset_desc, model_id, name=None):
        training_id = str(hashlib.md5(f'{str(dataset_desc)}{str(model_id)}'.encode('utf-8')).hexdigest())
        if name:
            training_id = f"{name}-{training_id}"
        return training_id

class Dataset:

    def __init__(self, fun=None, credentials_store=None):

        if not credentials_store:
            credentials_store = CredentialsStore()
        if not fun:
            fun = CloudFunctions(credentials_store)
        self.fun = fun

    def createDataset(self, pipeline, data, config, debug=False):

        if debug:
            print(f"Pipeline info - ID: {pipeline}, Config: {str(config)}")
            print(f"Sending data to pipeline : {str(data)}")

        data = {
            "data": json.loads(data.to_json(orient='split', date_format="iso")),
            "dataset": config
        }

        response = self.fun.call_function(pipeline, data, debug)

        if debug:
            print("Pipeline response : " + str(response))

        if response:
            dataset = pd.read_json(json.dumps(response), orient='split', dtype=False)
            if not dataset.empty:  # Check for empty dataset
                dataset.index.names = ["Date_Time"]  # Workaround for lack of index name in marshalling
                dataset = ppl.localize(dataset, "UTC", config["timezone"])  # Marshalling turns dates to UTC
            return dataset
        else:
            raise Exception("No response received from pipeline execution")

    @staticmethod
    def csvtojson(csv, dataset_desc, market, createId=True):
        obj = {}
        if (createId):
            obj["id"] = Dataset.generateId(dataset_desc, market)
        obj["dataset_desc"] = dataset_desc
        obj["market"] = market
        obj["data"] = csv.values.tolist()
        obj["tz"] = csv.index.tz.zone
        obj["index"] = [date.isoformat() for date in csv.index.tz_localize(None)]  # Remove locale
        return json.dumps(obj)

    @staticmethod
    def jsontocsv(jsonObj):
        return pd.DataFrame(jsonObj["data"], index=pd.DatetimeIndex(jsonObj["index"], name="Date_Time", tz=pytz.timezone(jsonObj["tz"])))


import telegram
import telegram.ext
import modules.chat_settings as chat_settings
import modules.strings
import modules.common as common
import modules.sql as sql
import modules.users_cache as users_cache

@common.private_only_dec()
async def cmd_help(update: telegram.Update, context: telegram.ext.CallbackContext):
    """/help"""
    await update.effective_chat.send_message(parse_mode=telegram.constants.ParseMode.HTML,
                                             text=modules.strings.CMD_HELP_TEXT)

@common.private_only_dec()
async def cmd_privacy(update: telegram.Update, context: telegram.ext.CallbackContext):
    """/privacy"""
    await update.effective_chat.send_message(parse_mode=telegram.constants.ParseMode.HTML,
                                             text=modules.strings.CMD_PRIVACY_TEXT)

async def generic_command_handler(update: telegram.Update, context: telegram.ext.CallbackContext):
    """Completely ignore channels and inline queries"""
    if update.effective_chat.type == telegram.constants.ChatType.CHANNEL or update.effective_chat.type == telegram.constants.ChatType.SENDER:
        raise telegram.ext.ApplicationHandlerStop

@common.group_only_dec()
async def cmd_bonk(update: telegram.Update, context: telegram.ext.CallbackContext):
    """/bonk"""
    common.ensure_chat_in_db(update.effective_chat)

    # Get settings for group
    settings = chat_settings.get_group_settings(update.effective_chat, settings=(
        chat_settings.show_help_msgs, chat_settings.announce_top_user, chat_settings.allow_counter_bonks))
    chat_show_help_msgs = settings[chat_settings.show_help_msgs.name]
    chat_announce_top_user = settings[chat_settings.announce_top_user.name]
    chat_allow_counter_bonks = settings[chat_settings.allow_counter_bonks.name]

    # Show help message if message isn't a reply and group is set to show help messages
    if update.effective_message.reply_to_message is None:
        if chat_show_help_msgs:
            await update.effective_chat.send_message(reply_to_message_id=update.effective_message.id,
                                                     text=modules.strings.NO_REPLIED_MESSAGE)
    else:
        # Get top score and holder to use later if announce new top user is enabled
        if chat_announce_top_user:
            # Since message IDs are incremental and per chat, we can use them to determine if a message is newer without storing timestamps
            top_score = sql.fetchone(f"SELECT COUNT(*) AS bonks_count, user_id, MAX(message_id) AS last_message "
                                     f"FROM bonks "
                                     f"WHERE chat_id = {update.effective_chat.id} "
                                     f"GROUP BY user_id "
                                     f"ORDER BY bonks_count DESC, last_message ASC "
                                     f"LIMIT 1")
        bonk_user_id = update.effective_message.reply_to_message.from_user.id
        bonk_message_id = update.effective_message.reply_to_message.message_id
        bonk_message = update.effective_message.reply_to_message
        block_counter_bonk = False
        already_bonked = False if sql.fetchone(f"SELECT * "
                                              f"FROM bonks "
                                              f"WHERE user_id = {bonk_user_id} "
                                              f"AND chat_id = {update.effective_chat.id} "
                                              f"AND message_id = {bonk_message_id}") is None else True

        # If counter bonks are disallowed, check if the message contains the /bonk command
        if not chat_allow_counter_bonks:
            dict = bonk_message.parse_entities([telegram.constants.MessageEntityType.BOT_COMMAND])
            for i in dict:
                if dict[i] == "/bonk" or dict[i] == f"/bonk{context.bot.name}":
                    block_counter_bonk = True
                    break

        # Block counter bonk
        if block_counter_bonk:
            if chat_show_help_msgs:
                await update.effective_chat.send_message(reply_to_message_id=update.effective_message.id,
                                                         text=modules.strings.COUNTER_BONKS_DISABLED)
        # Bot can't bonk itself
        elif update.effective_message.reply_to_message.from_user.id == context.bot.id:
            if chat_show_help_msgs:
                await update.effective_chat.send_message(reply_to_message_id=update.effective_message.id,
                                                         text=modules.strings.CANT_BONK_BOT)
        # Can't bonk the same message more than once
        elif already_bonked:
            if chat_show_help_msgs:
                await update.effective_chat.send_message(reply_to_message_id=update.effective_message.id,
                                                         text=modules.strings.ALREADY_BONKED_MESSAGE)
        else:
            # Insert bonk and update bonks count for current chat
            sql.run_query_s(f"INSERT INTO bonks "
                            f"VALUES ({bonk_user_id}, {update.effective_chat.id}, {bonk_message_id})")
            res = sql.fetchone(f"SELECT COUNT(*) "
                               f"FROM bonks "
                               f"WHERE chat_id = {update.effective_chat.id} "
                               f"AND user_id = {bonk_user_id}")
            message_text = modules.strings.replacements_BONK.replace("{?0?}", str(res)).replace("{?1?}", '' if res == 1 else 's')
            await update.effective_chat.send_message(reply_to_message_id=bonk_message_id,
                                                     text=message_text)
            # Announce new top user if that's the case
            if chat_announce_top_user and (top_score is None or (
                    res > top_score[0] and top_score[1] != update.effective_message.reply_to_message.from_user.id)):
                await update.effective_chat.send_message(reply_to_message_id=bonk_message_id,
                                                         text=modules.strings.HORNIEST_USER_ANNOUNCEMENT)

async def cmd_erase_data(update: telegram.Update, context: telegram.ext.CallbackContext):
    """/erase_data"""
    # Check if chat is group and user is admin
    if common.is_group(update.effective_chat):
        if await common.is_admin(update.effective_chat, update.effective_user.id):
            # Check if there's data stored about this group
            res = sql.fetchone(f"SELECT * "
                               f"FROM chats "
                               f"WHERE chat_id = {update.effective_chat.id}")
            if res is None:
                await update.effective_chat.send_message(text=modules.strings.NO_GROUP_DATA)
            else:
                await update.effective_chat.send_message(
                    text=modules.strings.ERASE_GROUP_WARNING,
                    reply_markup=modules.chat_settings.erase_data_inline_keyboard)
        elif chat_settings.get_group_settings(update.effective_chat, modules.chat_settings.allow_top_cmd):
            await update.effective_chat.send_message(text=modules.strings.ADMINS_ONLY,
                                                     reply_to_message_id=update.effective_message.id)
    else:
        # Check if there's data stored about this user
        res = sql.fetchone(f"SELECT * "
                           f"FROM bonks "
                           f"WHERE user_id = {update.effective_user.id}")
        if res is None:
            await update.effective_chat.send_message(text=modules.strings.NO_USER_DATA)
        else:
            await update.effective_chat.send_message(
                text=modules.strings.ERASE_USER_WARNING,
                reply_markup=modules.chat_settings.erase_data_inline_keyboard)

@common.group_only_dec()
async def cmd_top(update: telegram.Update, context: telegram.ext.CallbackContext):
    """/top"""
    # Check if chat is group and /top is allowed
    settings = modules.chat_settings.get_group_settings(update.effective_chat,
                                                        (chat_settings.allow_top_cmd, chat_settings.show_help_msgs))
    if settings[chat_settings.allow_top_cmd.name]:
        tot_chat_bonks = sql.fetchone(f"SELECT COUNT(*) "
                                      f"FROM bonks "
                                      f"WHERE chat_id = {update.effective_chat.id}")
        msg = f"Total number of bonks: {tot_chat_bonks}\n"
        # Get top 10 users
        top10 = sql.run_query_s(f"SELECT COUNT(*) AS bonks_count, user_id, MAX(message_id) AS last_message "
                                f"FROM bonks "
                                f"WHERE chat_id = {update.effective_chat.id} "
                                f"GROUP BY user_id "
                                f"ORDER BY bonks_count DESC, last_message ASC "
                                f"LIMIT 10").fetchall()
        # Check if there's data about this group
        if top10 is None:
            await update.effective_chat.send_message(reply_to_message_id=update.effective_message.id,
                                                     text=modules.strings.NO_BONKS_IN_CHAT)
        else:
            # Generate top 10
            for i in range(min(len(top10), 10)):
                try:
                    users_cache.update_user_info(await update.effective_chat.get_member(user_id=top10[i][1]))
                except telegram.error.BadRequest:
                    pass
                finally:
                    msg += f"{i + 1}. {users_cache.get_user_name(top10[i][1])}: {top10[i][0]}\n"
            await update.effective_chat.send_message(reply_to_message_id=update.effective_message.id,
                                                     text=msg,
                                                     parse_mode=telegram.constants.ParseMode.HTML,
                                                     disable_web_page_preview=True)
    elif settings[chat_settings.show_help_msgs.name]:
        await update.effective_chat.send_message(reply_to_message_id=update.effective_message.id,
                                                 text=modules.strings.COMMAND_DISABLED_BY_ADMINS)

@common.private_only_dec()
async def cmd_stats(update: telegram.Update, context: telegram.ext.CallbackContext):
    """/stats"""
    # Get stats
    global_bonks = common.sql.fetchone(f"SELECT COUNT(*) "
                                       f"FROM bonks "
                                       f"WHERE user_id = {update.effective_user.id}")
    most_group_bonks = common.sql.fetchone(f"SELECT COUNT(*) AS count "
                                           f"FROM bonks "
                                           f"WHERE user_id = {update.effective_user.id} "
                                           f"GROUP BY chat_id "
                                           f"ORDER BY count "
                                           f"LIMIT 1")
    if most_group_bonks is None:
        await update.effective_chat.send_message(text=modules.strings.NO_DATA_FOR_USER)
    else:
        await update.effective_chat.send_message(text=f"You were bonked a total of {global_bonks} time{'' if global_bonks == 1 else 's'}\n"
                                                      f"Your personal record is {most_group_bonks} time{'' if most_group_bonks == 1 else 's'} in a single chat") #it's 10pm i don't want to deal with this one rn

@common.private_only_dec()
async def cmd_start(update: telegram.Update, context: telegram.ext.CallbackContext):
    """/start"""
    await context.bot.send_message(chat_id=update.effective_chat.id,
                                   text=modules.strings.CMD_START)

@common.private_only_dec()
async def cmd_about(update: telegram.Update, context: telegram.ext.CallbackContext):
    """/about"""
    await context.bot.send_message(chat_id=update.effective_chat.id,
                                   parse_mode=telegram.constants.ParseMode.HTML,
                                   text=modules.strings.CMD_ABOUT_TEXT)


import sqlite3
import os
import matplotlib.pyplot as plt
import numpy as np
def openDatabase(db_name = 'CovidObesityProject'):
    path = os.path.dirname(os.path.abspath(__file__))
    conn = sqlite3.connect(path+'/'+db_name)
    cur = conn.cursor()
    return cur, conn
def getMainData(cur,con):

    avgData = []

    cur.execute(
            """
            SELECT CountryID.Country, PovertyRankData.PovertyRank,ObesityData.ObesityRate, CovidData.TotalCases,CovidData.TotalDeaths
            FROM CountryID
            JOIN CovidData ON CovidData.Country_ID = CountryID.Country_ID
            JOIN ObesityData on ObesityData.Country_ID=CovidData.Country_ID
            JOIN PovertyRankData on PovertyRankData.Country_ID=ObesityData.Country_ID
            """)
    data = cur.fetchall()
    return data

def deathRates(data):
    newList =[]
    for toop in data:
        if toop[3] == 0:
            rate = 0
        else:
            rate = round((toop[4]/toop[3]) * 100, 2)
        newT = (toop[0], rate)
        newList.append(newT)
    return newList

def writefile(data):
    path = os.path.dirname(os.path.abspath(__file__))
    with open(path + '/deathrates.txt', 'w') as f:
        f.write("Country, " + "Death Rates\n" )
        for toop in data:
            f.write(toop[0] + ", " + str(toop[1]) + '\n')
def ObesityDeathRateScatter(data):
    obesitylist=[]
    deathratelist=[]
    for country in data:
        if country[3] == 0:
            rate = 0
        else:
            rate = round((country[4]/country[3]) * 100, 2)
        obesitylist.append(country[2])
        deathratelist.append(rate)
    x = np.array(obesitylist)
    y = np.array(deathratelist)

    plt.scatter(x, y, color='red')
    plt.xlim([0, 40])
    plt.ylim([0,10])
    z = np.polyfit(x, y, 1)
    p = np.poly1d(z)
    plt.plot(x, p(x))
    plt.title('Obesity Rate and Death Rate Scatterplot')
    plt.xlabel('Obesity Rate (%)')
    plt.ylabel('Death Rate (%)')

    plt.show()

def PovertyDeathRateScatter(data):
    povertylist=[]
    deathratelist=[]
    for country in data:
        if country[3] == 0:
            rate = 0
        else:
            rate = round((country[4]/country[3]) * 100, 2)
        povertylist.append(country[1])
        deathratelist.append(rate)
    x = np.array(povertylist)
    y = np.array(deathratelist)

    plt.scatter(x, y, color='green')
    # plt.xlim([0, 40])
    plt.ylim([0,8])
    z = np.polyfit(x, y, 1)
    p = np.poly1d(z)
    plt.plot(x, p(x))
    plt.title('Poverty Rank and Death Rate Scatterplot')
    plt.xlabel('Poverty Rank')
    plt.ylabel('Death Rate (%)')

    plt.show()

def Top10DeathRate(data):
    tuplist=[]
    for country in data:
        if country[3] == 0:
            rate = 0
        else:
            rate = round((country[4]/country[3]) * 100, 2)
        tup=(country[0],rate)
        tuplist.append(tup)
    tuplist.sort(key=lambda y: y[1], reverse=True)
    #print(tuplist)
    tuplist10=tuplist[:10]

    namelist=[]
    ratelist=[]
    for x in tuplist10:
        namelist.append(x[0])
        ratelist.append(x[1])
    x=np.array(namelist)
    y=np.array(ratelist)

    plt.barh(x, y)
    plt.title('Top 10 Death Rate')
    plt.xlabel('Death Rate')
    plt.ylabel('Countries')

    plt.show()

def bottom10DeathRate(data):
    tuplist=[]
    for country in data:
        if country[3] == 0:
            rate = 0
        else:
            rate = round((country[4]/country[3]) * 100, 2)
        tup=(country[0],rate)
        tuplist.append(tup)
    tuplist.sort(key=lambda y: y[1])
    #print(tuplist)
    tuplist10=tuplist[:10]

    namelist=[]
    ratelist=[]
    for x in tuplist10:
        namelist.append(x[0])
        ratelist.append(x[1])
    x=np.array(namelist)
    y=np.array(ratelist)

    plt.barh(x, y, color='orange')
    plt.title('Bottom 10 Death Rate')
    plt.xlabel('Death Rate')
    plt.ylabel('Countries')

    plt.show()
def PovertyObesityScatter(data):
    povertylist=[]
    obesitylist=[]
    for country in data:
        povertylist.append(country[1])
        obesitylist.append(country[2])
    x = np.array(povertylist)
    y = np.array(obesitylist)

    plt.scatter(x, y, color='purple')
    # plt.xlim([0, 40])
    # plt.ylim([0,8])
    z = np.polyfit(x, y, 1)
    p = np.poly1d(z)
    plt.plot(x, p(x))
    plt.title('Poverty Rank and Obesity Rate Scatterplot')
    plt.xlabel('Poverty Rank (lower is poorer')
    plt.ylabel('Obesity Rate (%)')

    plt.show()
def Top10PovertyRate(data):
    tuplist=[]
    for country in data:
        tup=(country[0],country[1])
        tuplist.append(tup)
    tuplist.sort(key=lambda y: y[1])
    #print(tuplist)
    tuplist10=tuplist[:10]

    namelist=[]
    povertylist=[]
    for x in tuplist10:
        namelist.append(x[0])
        povertylist.append(x[1])
    x=np.array(namelist)
    y=np.array(povertylist)

    plt.barh(x, y, color="red")
    plt.title('Top 10 Poverty Rate')
    plt.xlabel('Poverty Rank(lower means poorer)')
    plt.ylabel('Countries')

    plt.show()

def main():
    c,d=openDatabase()
    da = getMainData(c,d)
    rat = deathRates(da)
    writefile(rat)
    ObesityDeathRateScatter(da)
    PovertyDeathRateScatter(da)
    Top10DeathRate(da)
    bottom10DeathRate(da)
    PovertyObesityScatter(da)
    Top10PovertyRate(da)
main()


import os
import sys

import numpy as np
from scipy.stats import ortho_group
import torch

ROOT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__),'..','..'))
if ROOT_DIR not in sys.path: sys.path.append(ROOT_DIR)

import response_contour_analysis.utils.principal_curvature as curve_utils

import pytest

DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'

class QuadraticFunction(torch.nn.Module):
    def __init__(self, diag=None, hess=None):
        super().__init__()

        if hess is not None:
            self.hessian = torch.tensor(hess).to(DEVICE)
        else:
            if diag is None:
                diag = [1.0, 2.0]
            self.hessian = torch.diag(torch.tensor(diag)).to(DEVICE)

        self.hessian = torch.nn.Parameter(self.hessian, requires_grad=False)

    def forward(self, x):
        return torch.dot(x, torch.matmul(self.hessian, x))

def value_grad_hess(f, point, dtype):
    value = f(point)
    grad = torch.autograd.functional.jacobian(f, point)
    hess = torch.autograd.functional.hessian(f, point)
    return value.to(dtype), grad.to(dtype), hess.to(dtype)

#def hyperboloid_eq(x, a, c):
#    return (x[:-1] ** 2) / (a ** 2) - x[-1] ** 2 / c ** 2

@pytest.mark.parametrize('curve_func', ['moosavi', 'golden', 'poole', 'lee_graph'])
def test_hyperboloid(curve_func):
    if curve_func == 'moosavi' or curve_func=='poole':
        pytest.skip() # we know this will fail
    a = 1
    c = 3
    f = lambda x: curve_utils.hyperboloid_graph(x[0], x[1], a=a, c=c)
    x = 2.0
    y = 2.0
    point = torch.tensor([x, y]).to(DEVICE)
    z, grad, hess = value_grad_hess(f, point, dtype=torch.double)

    shape_operator, principal_curvatures, principal_directions = curve_utils.local_response_curvature_alternates(grad, hess, curve_func)

    u = z.cpu().numpy() / c
    analytic_gauss_curvature = -c**2 / (c**2 + (a**2 + c**2) * u**2)**2

    np.testing.assert_allclose(
        analytic_gauss_curvature,
        torch.det(shape_operator).cpu().numpy(),
        rtol=1e-6
    )

@pytest.mark.parametrize('curve_func', ['moosavi', 'poole', 'lee_level'])
@pytest.mark.parametrize('dimensions', [2, 3, 5, 10])
@pytest.mark.parametrize('radius', [0.1, 1, 2, 3, 10])
def test_sphere_pc_variance(dimensions, radius, curve_func):
    if curve_func == 'moosavi' or curve_func=='poole':
        pytest.skip() # we know this will fail
    f = QuadraticFunction(np.ones(dimensions)).to(DEVICE)
    #point = torch.tensor(np.hstack((np.zeros(dimensions-1), [radius]))).to(DEVICE)
    point = torch.tensor(np.ones(dimensions) / np.sqrt(dimensions) * radius).to(DEVICE)
    value, pt_grad, pt_hess = value_grad_hess(f, point, dtype=torch.double)
    iso_shape_operator, iso_curvatures, iso_directions = curve_utils.local_response_curvature_alternates(pt_grad, pt_hess, curve_func)
    iso_curvatures = iso_curvatures.detach().cpu().numpy()
    assert iso_curvatures.var() < 1e-16

@pytest.mark.parametrize('curve_func', ['moosavi', 'poole', 'lee_level'])
@pytest.mark.parametrize('dimensions', [2, 3, 5, 10])
@pytest.mark.parametrize('radius', [0.1, 1, 2, 3, 10])
def test_sphere_pcs(dimensions, radius, curve_func):
    if curve_func == 'moosavi' or curve_func=='poole':
        pytest.skip() # we know this will fail
    f = QuadraticFunction(np.ones(dimensions)).to(DEVICE)
    #point = torch.tensor(np.hstack((np.zeros(dimensions-1), [radius]))).to(DEVICE)
    point = torch.tensor(np.ones(dimensions) / np.sqrt(dimensions) * radius).to(DEVICE)
    value, pt_grad, pt_hess = value_grad_hess(f, point, dtype=torch.double)
    iso_shape_operator, iso_curvatures, iso_directions = curve_utils.local_response_curvature_alternates(pt_grad, pt_hess, curve_func)
    iso_curvatures = iso_curvatures.detach().cpu().numpy()
    expected_principal_curvature = 1 / radius
    np.testing.assert_allclose(np.abs(iso_curvatures), [expected_principal_curvature,]*(dimensions-1))

@pytest.mark.parametrize('curve_func', ['moosavi', 'poole', 'lee_level'])
@pytest.mark.parametrize('dimensions', [2, 3, 5, 10])
@pytest.mark.parametrize('radius', [0.1, 1, 2, 3, 10])
def test_sphere_gauss(dimensions, radius, curve_func):
    if curve_func == 'moosavi' or curve_func=='poole':
        pytest.skip() # we know this will fail
    f = QuadraticFunction(np.ones(dimensions)).to(DEVICE)
    #point = torch.tensor(np.hstack((np.zeros(dimensions-1), [radius]))).to(DEVICE)
    point = torch.tensor(np.ones(dimensions) / np.sqrt(dimensions) * radius).to(DEVICE)
    value, pt_grad, pt_hess = value_grad_hess(f, point, dtype=torch.double)
    iso_shape_operator, iso_curvatures, iso_directions = curve_utils.local_response_curvature_alternates(pt_grad, pt_hess, curve_func)
    iso_curvatures = iso_curvatures.detach().cpu().numpy()
    expected_principal_curvature = 1 / radius
    expected_gaussian_curvature = np.prod([expected_principal_curvature,]*(dimensions-1))
    np.testing.assert_allclose(np.abs(np.prod(iso_curvatures)), expected_gaussian_curvature)

@pytest.mark.parametrize('dimensions', [2, 3, 5])
@pytest.mark.parametrize('radius', [0.1, 1, 2])
def test_coordinate_transform_sphere(dimensions, radius):
    f = QuadraticFunction(np.ones(dimensions)).to(DEVICE)
    #point = torch.tensor(np.hstack((np.zeros(dimensions-1), [radius]))).to(DEVICE)
    point = torch.tensor(np.ones(dimensions) / np.sqrt(dimensions) * radius).to(DEVICE)

    value, pt_grad, pt_hess = value_grad_hess(f, point, dtype=torch.double)

    coordinate_transformations = [
        None,
        -torch.eye(len(pt_grad), dtype=torch.double).to(DEVICE),
        torch.tensor(ortho_group.rvs(len(point)), dtype=torch.double).to(DEVICE)]
    #coordinate_transformation = torch.tensor(ortho_group.rvs(len(point)), dtype=torch.double).to(DEVICE)
    for coordinate_transformation in coordinate_transformations:
        iso_shape_operator, iso_curvatures, iso_directions = curve_utils.local_response_curvature_level_set(
            pt_grad, pt_hess, coordinate_transformation=coordinate_transformation)

        iso_curvatures = iso_curvatures.detach().cpu().numpy()

        assert iso_curvatures.var() < 1e-16

        expected_gaussian_curvature = 1 / (radius ** (dimensions - 1))
        np.testing.assert_allclose(np.abs(np.prod(iso_curvatures)), expected_gaussian_curvature)

@pytest.mark.skip("Will usually fail because the transformation reults in bad implicit functions")
def test_coordinate_transform():
    f = QuadraticFunction(np.array([1.0, 2.0, 0.1, 0.5])).to(DEVICE)
    point = torch.tensor(np.hstack((np.zeros(3), [1]))).to(DEVICE)

    value, pt_grad, pt_hess = value_grad_hess(f, point, dtype=torch.double)

    iso_shape_operator1, iso_curvatures1, iso_directions1 = curve_utils.local_response_curvature_level_set(pt_grad, pt_hess)
    iso_curvatures1 = iso_curvatures1.detach().cpu().numpy()

    M = ortho_group.rvs(len(point))
    #M = np.eye(len(point))
    #_M = ortho_group.rvs(len(point) - 1)
    #M[:-1, :-1] = _M
    coordinate_transformation = torch.tensor(M, dtype=torch.double).to(DEVICE)
    iso_shape_operator2, iso_curvatures2, iso_directions2 = curve_utils.local_response_curvature_level_set(pt_grad, pt_hess, coordinate_transformation=coordinate_transformation)
    iso_curvatures2 = iso_curvatures2.detach().cpu().numpy()

    np.testing.assert_allclose(iso_curvatures1, iso_curvatures2)

@pytest.mark.parametrize('dimensions', [20, 40, 60, 500])
@pytest.mark.parametrize('radius', [0.1, 1, 2])
def test_bad_condition(dimensions, radius):
    f_3d = QuadraticFunction([3.0, 3.0, 1.0]+[0.5] * (dimensions - 3)).to(DEVICE)
    point = torch.tensor(
        [np.sqrt(0.5), np.sqrt(0.5), 0]+[0] * (dimensions - 3),
        dtype=torch.float
    ).to(DEVICE) * radius

    # first in original coordinates
    value_orig, pt_grad_orig, pt_hess_orig = value_grad_hess(f_3d, point, dtype=torch.double)
    with pytest.raises(ValueError):
        iso_shape_operator_orig, iso_curvatures_orig, iso_directions_orig = \
            curve_utils.local_response_curvature_level_set(
                pt_grad_orig,
                pt_hess_orig,
                coordinate_transformation=torch.eye(len(point)).to(DEVICE)
            )

    rst = None
    # This will in general result in very bad conditioning
    M = ortho_group.rvs(len(point), random_state=rst)
    coordinate_transformation = torch.tensor(M).to(DEVICE)

    def new_f(x):
        return f_3d(torch.matmul(coordinate_transformation.T.type(x.dtype), x))
    new_point = torch.matmul(
        coordinate_transformation.type(point.dtype),
        point
    ).detach().clone()

    value_new, pt_grad_new, pt_hess_new = value_grad_hess(new_f, new_point, dtype=torch.double)

    np.testing.assert_allclose(
        value_orig.cpu().numpy(),
        value_new.cpu().numpy(),
        rtol=1e-6
    )

    iso_shape_operator_new, iso_curvatures_new, iso_directions_new = \
        curve_utils.local_response_curvature_level_set(
            pt_grad_new, pt_hess_new,
        )

    expected_curvatures = np.array(
        [-1/6] * (dimensions - 3) + [-1/3, -1.0]
    ) / radius

    np.testing.assert_allclose(
        iso_curvatures_new.cpu().numpy(),
        expected_curvatures,
        rtol=1e-6,
    )