from bs4 import BeautifulSoup
import requests
import requests_html
from requests_html import HTMLSession
import time
from typing import Optional, List, Dict
import urllib

RESPONSE = requests.models.Response
URL = str

HEADERS = [
    None,
    {
        'User-Agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_11_5) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/50.0.2661.102 Safari/537.36'
    },
    {
        'Accept-Encoding': 'gzip, deflate, sdch',
        'Accept-Language': 'en-US,en;q=0.8',
        'Upgrade-Insecure-Requests': '1',
        'User-Agent': 'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/56.0.2924.87 Safari/537.36',
        'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8',
        'Cache-Control': 'max-age=0',
        'Connection': 'keep-alive',
    },
]

class GGScrapException(Exception):
    pass

# waits for a connection, returns True if connected
def wait_for_connection(
        url: URL=               'http://ident.me/',
        logger=                 None,
        max_sec: Optional[int]= None,
) -> bool:

    connected = False
    loop_ix = 0
    while True:
        response = None
        try:
            response = requests.get(url=url, timeout=5)
            if logger: logger.info(f'waiting for connection, got: {response}')
        except Exception as e:
            if logger: logger.warning(f'got connection exception: {e}')
        if response and response.status_code == 200:
            connected = True
            if logger: logger.info(f'got nice connection!')
            break
        time.sleep(5)

        loop_ix += 1
        if max_sec is not None and max_sec <= loop_ix * 5:
            break

    return connected

# tries to download RESPONSE from URL
def download_response(
        url: URL,
        header: Optional[Dict]= None,
        proxy: Optional[str]=   None,
        logger=                 None) -> Optional[RESPONSE]:
    try:
        session = HTMLSession()
        proxies = {'http': f'http://{proxy}'} if proxy else None
        response = session.get(url, headers=header, proxies=proxies)
        if logger: logger.debug(f'download_response() got response: "{response}" from url: {url}, header: {header}, proxy: {proxy}')
        response.session = session
        return response
    except Exception as e:
        msg = f'download_response() got exception: "{e}", url: {url}, header: {header}, proxy: {proxy}'
        if logger: logger.warning(msg)
        return None

# extracts sub-urls from RESPONSE
def extract_subURLs(response:RESPONSE, logger=None) -> List[URL]:
    try:
        html = requests_html.HTML(
            session=            HTMLSession(),
            url=                response.url,
            html=               response.content,
            default_encoding=   response.encoding)
        return list(html.absolute_links)
    except Exception as e:
        msg = f'extract_subURLs() got exception: "{e}"'
        if logger: logger.warning(msg)
        return []

# extracts words from the url base
def get_head_words(url:URL) -> List[str]:
    url = url.replace('https://','')
    url = url.replace('http://', '')
    head = url.split('/')[0]
    head = head.split('.')
    head.pop(-1) # remove domain
    if 'www' in head: head.remove('www')
    return head

# filters given URLs, removes those containing any filter before first /
def filter_URLs(
        urls: List[URL],
        filters: List[str] # filter is a str, words space separated, like 'maps google', if both are in URL before / <- filter matches and URL is filtered out
) -> List[URL]:

    filters_split = [f.split() for f in filters]
    filtered = []
    for url in urls:
        head_words = get_head_words(url)
        url_ok = True
        for fs in filters_split:
            all_in = True
            for f in fs:
                if f not in head_words:
                    all_in = False
                    break
            if all_in:
                url_ok = False
                break
        if url_ok: filtered.append(url)

    return filtered

# extracts text from RESPONSE
def extract_text(
        response: RESPONSE,
        separator=      '\n',
        encode_decode=  True,
        logger=         None,
) -> Optional[str]:
    try:
        soup = BeautifulSoup(response.text, 'html.parser')
        text = soup.get_text(separator=separator)
        if encode_decode:
            text = text.encode("ascii", "ignore")
            text = text.decode()
        return text
    except Exception as e:
        msg = f'get_texts() got exception: "{e}"'
        if logger: logger.warning(msg)
        return None

# builds google search URL for given query
def build_google_url(query:str) -> str:
    if not query:
        raise GGScrapException('cannot build google url for empty query')
    return f'https://www.google.com/search?q={urllib.parse.quote_plus(query)}'


#model utilities
from sklearn.utils import shuffle
from sklearn.model_selection import KFold
from sklearn.metrics import accuracy_score

def convert_lables(Y_data):
  for i in range(len(Y_data)):
    if Y_data[i] == 0 : Y_data[i] = -1
  return Y_data

def train_with_crossvalidation(splits, X_data, Y_data, model):

  kf = KFold(n_splits = splits)
  kf.get_n_splits(X_data)
  results = []

  for train_index, test_index in kf.split(X_data): #indexes of splitted data

      X_train, X_test = X_data[train_index], X_data[test_index]
      Y_train, Y_test = Y_data[train_index], Y_data[test_index]

      X_train, Y_train = shuffle(X_train, Y_train)
      model = model.fit(X_train, Y_train)
      predictions = model.predict(X_test)
      results.append(accuracy_score(Y_test, predictions))

  return model, results


'''
very hacky first part answer to this day's problem,
it actually tracks which elf as well as how many calories that elf has.
'''
file = "input-official.txt"

# use this to track the answers
def update(total, idx):
    global highest, elf
    if highest < total:
        highest = total
        elf = idx

with open(file=file) as f:

    # track everything where the index will be the elf number (minus one)
    highest, total, idx, elf =  0, 0, 0, 0

    for line in f:

        if line != "\n":
            total+= int(line)
            print("adding on amount")

        else:
            print("newline detected")
            if total:
                idx +=1

            # add last one onto the tracker

            update(total, idx)

            # reset the running total for next line
            total = 0
    else:
        # python for/else construct -- in case it's end of file now.
        idx +=1
        update(total, idx)

print(highest, "amount")
print(elf, "which elf")


import array

def print_bin(x):
    for i in range(0, 8):
        if x & 128 > 0:
            print('1', end='')
        else:
            print('0', end='')

        x <<= 1

    print('')

def count_bits(x):
    num_bits = 0

    while x:
        num_bits += x & 1
        x >>= 1

    return num_bits

def parity(x):
    result = 0

    while x:
        result ^= x & 1

        x >>= 1

    return result

def pow2(x):
    count = 0

    while x:
        count+=1
        if count>1:
            return False

        x &= x - 1
    return True

def parity_set_bit(x):
    parity = 0

    while x:
        parity ^= 1 # flip the parity

        x &= x - 1

    return parity

def parity_lookups(x):
    mask_size = 16
    bit_mask = 0xFFFF

    return lookup[x >> (3 * mask_size)] ^ lookup[x >> (2 * mask_size) & bit_mask] ^ lookup[x >> mask_size & bit_mask] ^ lookup[x & bit_mask]

def parity_word(x):
    x ^= x >> 32
    x ^= x >> 16
    x ^= x >> 8
    x ^= x >> 4
    x ^= x >> 2
    x ^= x >> 1

    return x & 1

def parity_word_lookup(x):
    bit_mask = 0xFFFF

    x ^= x >> 32
    x ^= x >> 16

    return lookup[x & bit_mask]

lookup = array.array('b');

for i in range(0, 65536):
    lookup.append(parity_set_bit(i))

mask = 63
x = 77

print(bin(mask))
print(bin(x))
print(x & mask)

#print(count_bits(101))
# print(parity(3))
# print(parity_store(3))
# print(parity_set_bit(3))
# print(bin(65))
# print(bin(128))

print(parity(573121423114111))
print(parity_set_bit(573121423114111))
print(parity_lookups(573121423114111))
print(parity_word(573121423114111))
print(parity_word_lookup(573121423114111))

# print(bin(86))
# print_bin(86)

#print(lookup[554])

x = int('01010000', 2)
right = x & ~(x-1)
right -= 1
print(bin(x ^ right))

# 1010110
# 1010101 (x-1)
# 0101010 ~(x-1)
# 0000010 x & ~(x-1)

print(pow2(2))
print(pow2(3))
print(pow2(4))
print(pow2(6))
print(pow2(10))
print(pow2(16))


import os
import requests as res

print("web-scraper  v0.1")

#asking for the url
url = input("enter the url you want to scrape(give file url, if you give the direct url it will download the index page.):")

#asking for name of the file
name = input("enter the file name:")

#asking for the directory
dir = input("enter the directory")

#asking for thw user to wait
print("wait..")

#downloading the content
rq = res.get(url)

#making a directory
os.mkdir(dir)

#making a statement for directory at the end
dirst = "check your files in " + dir

#making the dir for making file
dir = dir + "/" + name

#making file
with open(dir , 'w') as f:
    f.write(rq.text)

#print the final statements
print("done!")
print(dirst)


a = [1,2,1,2,1,2,1,2,1]
print(a)

#Append
a.append(6)
print(a)

b = ["mango","apple","grapes"]
a.append(b)

print(a)

#Extend

a = [1,1,1,5,2,4,2,2,8,2,4,5,5,55]
print(a)
b = [4,4,4,4,4,5,4,4]
a.extend(b)
print(a)

#Index
a.index(5)

#Insert
a.insert(0,"aman")
print(a)

#remove
a.remove(0)
print(a)

#reverse
x = [45,4,5,4,4,5,4,5,4,5]
x.reverse()
print(x)

#count
x.count(4)

Chapter 20 List/List methods and supported operators.py


#Известно, что рост футболистов в сборной распределен нормально с дисперсией генеральной совокупности,
#равной 25 кв.см. Объем выборки равен 27, среднее выборочное составляет 174.2.
#Найдите доверительный интервал для математического ожидания с надежностью 0.95.
import math
from scipy.stats import norm

mean = 174.2
variance = 25
n = 27
alpha = 0.05

se = math.sqrt(variance) / math.sqrt(n)
z = norm.ppf(1 - alpha / 2)

lower = mean - z * se
upper = mean + z * se

print("Доверительный интервал: ({:.2f}, {:.2f})".format(lower, upper))


# only downloader tasks are done in the functions below

from urllib import request
from urllib.request import Request, urlopen # for get_article_webpage_downloader

def get_page_downloader():
    url = "https://outbreaks.globalincidentmap.com/home.php"
    return request.urlopen(url).read()

# using the difference ID list, append the ID to url and return
# a list of html text dumps containing individual article information
def get_details_list_downloader(difference_list):
    url = "https://outbreaks.globalincidentmap.com/eventdetail.php?ID="
    html_dump_list = []
    print("Found " + str(len(difference_list)) + " pages to crawl")
    count=1
    for ID in difference_list:
        curtr_url = url + str(ID)
        # put a try except block or check if empty string is returned here
        # attatch ID to the html dump, ID used in the data
        html_dump_list.append([request.urlopen(curtr_url).read(), ID])

        print("Downloaded " + str(count) + " of " + str(len(difference_list)))
        count=count+1
    return html_dump_list

# the article websites will kick scrapers, need to fool the security features
# and pretend the request is a request from a user-agent with a browser
# common security feature to avoid bots and scrapersself.
# this is indicated by a 403 error meaning understood request but not fulfill
def get_article_webpage_downloader(url):
    # adds headers to HTTP request to fool server
    try:
        req = Request(url, headers={'User-Agent': 'Mozilla/5.0'})
        x = urlopen(req, timeout=20).read()
        return x
    except:
        print("Failed to open url " + url)
        return ""


from language.python import ast
from analysis.storegraph import storegraph
from ... import intrinsics

import collections

def reindexUnionFind(uf):
	reindexed = {}

	for name in uf:
		group = uf[name]
		if group not in reindexed:
			reindexed[group] = [name]
		else:
			reindexed[group].append(name)

	return reindexed

class Renamer(object):
	def __init__(self):
		self.names = collections.defaultdict(lambda: -1)

	def __call__(self, name):
		self.names[name] += 1
		uid = self.names[name]

		if uid == 0:
			return name
		else:
			return self("%s_%d" % (name, uid))

class SamplerGroup(object):
	def __init__(self, canonical, group, uid):
		self.t = canonical.xtype.obj.pythonType()
		self.canonical = canonical
		self.group = group
		self.name = "samplerGroup%d" % uid
		self.impl = None

		assert len(group) == 1

		self.unique = True

LOCAL   = 'LOCAL'
UNIFORM = 'UNIFORM'
INPUT   = 'INPUT'
OUTPUT  = 'OUTPUT'

class Mergable(object):
	def __init__(self):
		self._forward = None

	def forward(self):
		if self._forward:
			self._forward = self._forward.forward()
			assert self._forward._forward is None
			return self._forward
		else:
			return self

class SubrefInfo(Mergable):
	def __init__(self):
		Mergable.__init__(self)
		self.name = None
		self.slots = set()
		self.holdsVolatile = False
		self.builtin = False
		self.impl = None

	def merge(self, analysis, other):
		assert isinstance(other, SubrefInfo), other

		self  = self.forward()
		other = other.forward()

		if self is other:
			return self

		other._forward = self

		self.holdsVolatile |= other.holdsVolatile
		self.slots.update(other.slots)

		self.mergeData(other)

		return self

	def postProcess(self, exgraph):
		self.findUnique(exgraph)

	def findUnique(self, exgraph):
		unique = True

		for ref in self.refs:
			unique &= ref.annotation.unique

		if unique:
			exclusive = exgraph.mutuallyExclusive(*self.refs)
			unique &= exclusive

		self.unique = unique

	def intrinsic(self):
		return False

	def setBaseName(self, base, nameCallback):
		if self.name is None:
			if self.holdsVolatile: base = "pool"
			self.name = nameCallback("%s_%s" % (base, self.postfix()))

	def markHoldsVolitile(self, analysis):
		self.holdsVolatile = True
		self.markRefsVolatile(analysis)

class TypeInfo(SubrefInfo):
	@property
	def t(self):
		return int

	def postfix(self):
		return "type"

	def accept(self, visitor, ref):
		return visitor.visitType(ref, self)

class IntrinsicSubrefInfo(SubrefInfo):
	def __init__(self, t):
		SubrefInfo.__init__(self)
		self.t = t
		self.sampler = t in intrinsics.samplerTypes
		self.refs = set()
		self.impl = None

	def postfix(self):
		return self.t.__name__

	def addRef(self, ref):
		if ref not in self.refs:
			self.refs.add(ref)

	def mergeData(self, other):
		assert self.t is other.t
		self.refs.update(other.refs)

	def intrinsic(self):
		return True

	def accept(self, visitor, ref):
		if self.holdsVolatile:
			if self.unique:
				return visitor.visitIntrinsicSingleton(ref, self)
			else:
				return visitor.visitIntrinsicPool(ref, self)
		else:
			return visitor.visitIntrinsicValue(ref, self)

	def markRefsVolatile(self, analysis):
		for ref in self.refs:
			objInfo = analysis.objInfo(ref)
			objInfo.markVolatile(analysis)

class SamplerSubrefInfo(IntrinsicSubrefInfo):
	def accept(self, visitor, ref):
		if self.holdsVolatile:
			assert False, "samplers should not be volatile?"
		else:
			return visitor.visitSampler(ref, self)

class SubpoolLUT(object):
	def __init__(self):
		self.type = None
		self.subpools = {}

	def typeSubpool(self):
		if self.type is None:
			self.type = TypeInfo()
			self.subpools['type'] = self.type
		return self.type

	def getName(self, name):
		sub = self.subpools[name]
		forward = sub.forward()

		if sub is not forward:
			self.subpools[name] = forward
			sub = forward
		return sub

	def intrinsic(self, t):
		assert t in intrinsics.intrinsicTypes, t
		if t not in self.subpools:
			sub = IntrinsicSubrefInfo(t)
			self.subpools[t] = sub
		else:
			sub = self.getName(t)
		return sub

	def sampler(self, t):
		assert t in intrinsics.samplerTypes, t
		if t not in self.subpools:
			sub = SamplerSubrefInfo(t)
			self.subpools[t] = sub
		else:
			sub = self.getName(t)
		return sub

# TODO this is what gets dirty?
class ObjectInfo(object):
	def __init__(self, name):
		self.name = name
		self.subrefs = set()
		self.volatile = False

	def markVolatile(self, analysis):
		if not self.volatile:
			self.volatile = True
			analysis.dirty.add(self)

	def resolve(self, analysis):
			sub = self.subrefs.pop().forward()
			for other in self.subrefs:
				sub = sub.merge(analysis, other)

			self.subrefs.clear()
			self.subrefs.add(sub)

			sub.markHoldsVolitile(analysis)

class ReferenceInfo(object):
	def __init__(self, output=False):
		self.name = None

		self.slots = set()
		self.refs  = set()

		self.lut = SubpoolLUT()

		self.final = True
		self.output = output

		self.anchor = False

		self.next  = set()
		self.prev  = set()

		self.fieldGroupsRefer = []
		self.containedFieldGroups = []

		self._forward = None

		self.dirty = False

		self.mode = LOCAL

	def subpools(self):
		return self.lut.subpools.itervalues()

	def addSlot(self, slot):
		self.slots.add(slot)

	def addRef(self, ref):
		self.refs.add(ref)

		t = ref.xtype.obj.pythonType()
		if t in intrinsics.samplerTypes:
			sub = self.lut.sampler(t)
			sub.output = self.output
			sub.slots.update(self.slots)
			sub.addRef(ref)
		elif intrinsics.isIntrinsicType(t):
			sub = self.lut.intrinsic(t)
			sub.output = self.output
			sub.slots.update(self.slots)
			sub.addRef(ref)
		else:
			# TODO fields?
			sub = None

		self.final &= ref.annotation.final

		return sub

	def forward(self):
		if self._forward:
			self._forward = self._forward.forward()
			assert self._forward._forward is None
			return self._forward
		else:
			return self

	def transfer(self, analysis, other):
		if other is not self:
			self.next.add(other)
			other.prev.add(self)

		return self

	def merge(self, analysis, other):
		self  = self.forward()
		other = other.forward()

		if self is other:
			return self

		assert not self.anchor and not other.anchor

		other._forward = self

		self.slots.update(other.slots)
		self.refs.update(other.refs)

		self.final &= other.final

		# Note that prev and next stay unforwarded, as we'll need to forward
		# everything when we use it, anyways
		self.prev.update(other.prev)
		self.next.update(other.next)

		if other.dirty:
			analysis.dirty.remove(other)

		if not self.dirty:
			analysis.markDirty(self)

		return self

	def updateLinks(self):
		self.prev = set([other.forward() for other in self.prev])
		self.next = set([other.forward() for other in self.next])

		if self in self.prev:
			self.prev.remove(self)

		if self in self.next:
			self.next.remove(self)

	def contract(self, analysis):
		self = self.forward()
		self.updateLinks()

		if self.anchor: return

		others = self.prev.union(self.next)
		for other in others:
			other = other.forward()
			if other.anchor: continue

#			if self.final == other.final:
#				print "CONTRACT"
#				print self.slots
#				print other.slots
#				print
#
#				self = self.merge(analysis, other)

	def postProcess(self):
		# Choose a random name
		self.name = tuple(self.slots)[0]

		self.types = set([ref.xtype.obj.pythonType() for ref in self.refs])

		# Make sure it exists
		if self.multipleTypes(): self.lut.typeSubpool()

	def multipleTypes(self):
		return len(self.types) > 1

	def setName(self, name, uniqueCallback):
		self.name = name

		for sub in self.subpools():
			sub.setBaseName(name, uniqueCallback)

	def copyNames(self, other):
		self.name = other.name
		for group, sub in self.lut.subpools.iteritems():
			othersub = other.lut.subpools[group]
			sub.name = othersub.name

	def setSpecialName(self, name):
		self.name = name

		for sub in self.subpools():
			sub.name = name
			sub.builtin = True

	def isSingleton(self):
		return self.unique

	def isStructure(self):
		return not self.unique and self.final

	def isPool(self):
		return not self.unique and not self.final


import tkinter
from quiz_brain import QuizBrain

THEME_COLOR = "#375362"

class QuizInterface:
    def __init__(self, quiz_brain: QuizBrain):
        self.quiz = quiz_brain
        self.window = tkinter.Tk()
        self.window.title("Quizzler")
        self.window.config(padx=20, pady=20, bg=THEME_COLOR)

        self.score_label = tkinter.Label()
        self.score_label.config(text=f"Score: 0/{self.quiz.total_questions}", fg="white", font=("Segoe UI", 12, "bold"),
                                bg=THEME_COLOR)
        self.score_label.grid(column=1, row=0)

        self.canvas = tkinter.Canvas(bg="white", width=300, height=250, highlightthickness=0)
        self.canvas.grid(column=0, row=1, columnspan=2, pady=50)
        self.question_text = self.canvas.create_text(150, 125, text="Question goes here.", font=("Arial", 18, "italic"),
                                                     width=280)

        check_button_path = tkinter.PhotoImage(file="./images/true.png")
        x_button_path = tkinter.PhotoImage(file="./images/false.png")
        self.check_button = tkinter.Button()
        self.check_button.config(borderwidth=0, highlightthickness=0, image=check_button_path,
                                 command=self.check_pressed)
        self.check_button.grid(column=0, row=2)
        self.x_button = tkinter.Button()
        self.x_button.config(borderwidth=0, highlightthickness=0, image=x_button_path, command=self.x_pressed)
        self.x_button.grid(column=1, row=2)

        self.get_next_question()

        self.window.mainloop()

    def get_next_question(self):
        self.canvas.config(bg="white")
        self.check_button.config(state="active")
        self.x_button.config(state="active")
        if self.quiz.still_has_questions():
            q_text = self.quiz.next_question()
            self.score_label.config(text=f"Score: {self.quiz.score}/{self.quiz.total_questions}")
            self.canvas.itemconfig(self.question_text, text=q_text)
        else:
            self.canvas.itemconfig(self.question_text, text="You've reached the end of the quiz")
            self.check_button.config(state="disabled")
            self.x_button.config(state="disabled")

    def check_pressed(self):
        self.give_feedback(self.quiz.check_answer("True"))

    def x_pressed(self):
        self.give_feedback(self.quiz.check_answer("False"))

    def give_feedback(self, is_right: bool):
        if is_right is True:
            self.canvas.config(bg="green")
        else:
            self.canvas.config(bg="red")
        self.check_button.config(state="disabled")
        self.x_button.config(state="disabled")
        self.window.after(1000, func=self.get_next_question)

Seção 34; Day 34 - Intermediate+ API Practice - Creating a GUI Quiz App/303. Trivia Question API Challenge/ui.py


def DisplayMulti(*VarArgs):
    for Arg in VarArgs:
        if Arg.upper() == 'CONT':
            continue
            print('Continue Argument: ' + Arg)
        elif Arg.upper() == 'BREAK':
            break
            print('Break Argument: '+ Arg)
        print('Good Argument '+ Arg)
DisplayMulti("a", "b", "cont","c")


import pyttsx3
import PyPDF2

engine = pyttsx3.init("sapi5")
book = open('sampleaudiobook.pdf','rb')
pdfReader = PyPDF2.PdfFileReader(book)
pages = pdfReader.numPages
print("The total number of pages in pdf is", pages)
speaker = pyttsx3
page = pdfReader.getPage(0)     #list range starts from 0
text = page.extractText()
engine.say(text)
print("But it reads only one page Because the index is 0.")
#engine.say('Just a sample text')
engine.runAndWait()


from pathlib import Path

import gspread

from scripts._dir import PARSER_OOT

SPREADSHEET_ID = "1SSFQgfg-EFqfRzdKnoyRZIB7t6cwc_4MyseFXRdOSqs"

gc: gspread.Client | None = None
workbook: gspread.Spreadsheet | None = None

def get_worksheet(name: str):
    global gc, workbook
    if gc is None:
        gc = gspread.oauth(
            credentials_filename=PARSER_OOT / "secrets/google-credentials.json",
            authorized_user_filename=PARSER_OOT / "secrets/google-token.json",
        )
    if workbook is None:
        workbook = gc.open_by_key(SPREADSHEET_ID)

    try:
        return workbook.worksheet(name)
    except gspread.WorksheetNotFound:
        return workbook.add_worksheet(name, 1, 1)


import numpy as np
import torch
import torch.optim as optim
import torch.nn.functional as F
import torch.utils.data as Data
import random
import pandas as pd
from utils_graph import *
from models_graph import GNNs
import argparse
import dgl
import os

parser = argparse.ArgumentParser()
parser.add_argument('--task_type', type=str, default='Regression',
                    choices=['Classification','Regression'])
parser.add_argument('--seed', type=int, default=1, help='Random seed.')
parser.add_argument('--epochs', type=int, default=100,
                    help='Number of epochs to train.')
parser.add_argument('--lr', type=float, default=0.005,
                    help='Initial learning rate.')
parser.add_argument('--hidden', type=int, default=256,
                    help='Number of hidden units.')
parser.add_argument('--src_vocab_size', type=int, default=15) # number of beads
parser.add_argument('--embed_dim', type=int, default=256)
parser.add_argument('--batch_size', type=int, default=1024)
parser.add_argument('--model', type=str, default='GraphSAGE',choices=['GCN','GAT','GraphSAGE'])
parser.add_argument('--GraphSAGE_aggregator', type=str, default='lstm',choices=['mean', 'gcn', 'pool', 'lstm'])
parser.add_argument('--GAT_heads', type=int, default=4)
parser.add_argument('--model_path', type=str, default='model_val_best.pt')

args = parser.parse_args()

args.model_path='{}_lr_{}_hid_{}_bs_{}.pt'.format(args.model,args.lr,args.hidden,args.batch_size)

device = 'cuda' if torch.cuda.is_available() else 'cpu'
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if torch.cuda.is_available(): torch.cuda.manual_seed(args.seed)

def main():

    if args.task_type == 'Classification':
        train_dataset, train_label = dgl.load_graphs('Graphical_Peptides/dgl_train_graphs_cla.bin')
        valid_dataset, valid_label = dgl.load_graphs('Graphical_Peptides/dgl_valid_graphs_cla.bin')
        test_dataset, test_label = dgl.load_graphs('Graphical_Peptides/dgl_test_graphs_cla.bin')
        train_label = train_label['labels'].long()
        valid_label = valid_label['labels'].long()
        test_label = test_label['labels'].long()
    elif args.task_type == 'Regression':
        train_dataset, train_label = dgl.load_graphs('Graphical_Peptides/dgl_train_graphs.bin')
        valid_dataset, valid_label = dgl.load_graphs('Graphical_Peptides/dgl_valid_graphs.bin')
        test_dataset, test_label = dgl.load_graphs('Graphical_Peptides/dgl_test_graphs.bin')
        train_label = train_label['labels'].float()
        valid_label = valid_label['labels'].float()
        test_label = test_label['labels'].float()

    dataset = MyDataSet(train_dataset,train_label)
    train_loader = Data.DataLoader(dataset, args.batch_size, True,
                         collate_fn=collate)
    valid_set = MyDataSet(valid_dataset,valid_label)
    val_graphs,val_labels = collate(valid_set)

    test_set = MyDataSet(test_dataset,test_label)
    test_graphs,test_labels = collate(test_set)

    valid_mse_saved = 100
    valid_acc_saved = 0
    loss_mse = torch.nn.MSELoss()
    # loss_cla = torch.nn.BCELoss()

    if args.task_type == 'Classification':
        model = GNNs(args, 2).to(device)
        optimizer = optim.Adam(model.parameters(), lr=1e-4)
    elif args.task_type == 'Regression':
        model = GNNs(args, 1).to(device)
        optimizer = optim.Adam(model.parameters(), lr=args.lr)

    for epoch in range(args.epochs):
        model.train()
        for graphs,labels in train_loader:
            graphs = graphs.to(device)
            labels = labels.to(device)
            outputs = model(graphs)

            if args.task_type == 'Classification':
                loss = F.nll_loss(outputs, labels)
                # loss = loss_cla(outputs, labels)
            elif args.task_type == 'Regression':
                loss = loss_mse(outputs, labels)

            # print('Epoch:',epoch,', Training Loss:',loss)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        if (epoch+1) % 1 == 0:
            model.eval()
            predict = model(val_graphs)

            if args.task_type == 'Classification':
                valid_acc = accuracy(predict, val_labels)
                if valid_acc_saved < valid_acc:
                    valid_acc_saved = valid_acc
                    print('Epoch:',epoch+1)
                    print('Valid Performance:',valid_acc.item())
                    torch.save(model.state_dict(),args.model_path)
                    print('Test Performance:',accuracy(model(test_graphs),test_labels).item())

            elif args.task_type == 'Regression':
                valid_mse = loss_mse(predict,val_labels).item()
                if valid_mse_saved > valid_mse:
                    valid_mse_saved = valid_mse
                    print('Epoch:',epoch+1)
                    print('Valid Performance:',valid_mse)
                    torch.save(model.state_dict(),args.model_path)
                    print('Test Performance:',loss_mse(model(test_graphs),test_labels).item())

    if args.task_type == 'Classification':
        predict = []
        model_load = GNNs(args, 2).to(device)
        checkpoint = torch.load(args.model_path)
        model_load.load_state_dict(checkpoint)
        model_load.eval()

        outputs = model_load(test_graphs)
        predict_test = outputs.max(1)[1].type_as(test_label)
        predict = predict + predict_test.cpu().detach().numpy().tolist()

        df_test_save = pd.DataFrame()
        labels = test_label.squeeze(1).tolist()
        df_test_seq = pd.read_csv('Sequential_Peptides/test_seqs_cla.csv')
        df_test_save['feature'] = df_test_seq['Feature']
        df_test_save['predict'] = predict
        df_test_save['label'] = labels
        test_acc = accuracy(model_load(test_graphs),test_labels).item()
        df_test_save['Acc'] = test_acc

        from sklearn.metrics import precision_score, recall_score, f1_score
        df_test_save['Precision'] = precision_score(test_label.cpu(),predict_test.cpu().detach())
        df_test_save['Recall'] = recall_score(test_label.cpu(),predict_test.cpu().detach())
        df_test_save['F1-score'] = f1_score(test_label.cpu(),predict_test.cpu().detach())

        df_test_save.to_csv('results_graph/Test_cla_{}_Acc_{}_lr_{}_hid_{}_bs_{}.csv'.format(args.model,test_acc,args.lr,args.hidden,args.batch_size))

    if args.task_type == 'Regression':
        predict = []
        model_load = GNNs(args, 1).to(device)
        checkpoint = torch.load(args.model_path)
        model_load.load_state_dict(checkpoint)
        model_load.eval()

        outputs = model_load(test_graphs)
        predict = predict + outputs.squeeze(1).cpu().detach().numpy().tolist()

        df_test_save = pd.DataFrame()
        labels = test_label.squeeze(1).tolist()
        df_test_seq = pd.read_csv('Sequential_Peptides/test_seqs.csv')
        df_test_save['feature'] = df_test_seq['Feature']
        df_test_save['predict'] = predict
        df_test_save['label'] = labels
        error = []
        for i in range(len(labels)):
            error.append(labels[i]-predict[i])
        absError = []
        squaredError = []
        for val in error:
            absError.append(abs(val))
            squaredError.append(val*val)
        MSE = sum(squaredError)/len(squaredError)
        MAE = sum(absError)/len(absError)

        from sklearn.metrics import r2_score
        R2 = r2_score(test_label.cpu(),outputs.cpu().detach())

        df_test_save['MSE'] = squaredError
        df_test_save['MAE'] = absError
        df_test_save['MSE_ave'] = MSE
        df_test_save['MAE_ave'] = MAE
        df_test_save['R2'] = R2
        df_test_save.to_csv('results_graph/Test_reg_{}_MAE_{}_lr_{}_hid_{}_bs_{}.csv'.format(args.model,MAE,args.lr,args.hidden,args.batch_size))

    os.remove(args.model_path)

if __name__ == '__main__':
    main()


# This script runs through the extracted audio lines in a folder
# and reproduces them one by one. Upon hearing it, you can enter
# the information you know about it as prompted. This metadata
# will be saved to a .csv file for indexing.

import os
import csv
import subprocess

def main():
    check_paths()
    handle__tracks_info_file()

def check_paths():
    # Make sure important paths and dirs are created

    #region # General paths
    global SCRIPT_PATH # Where this script is
    global SCRIPT_DIR_PATH # This script's dir
    global PROJECTS_DB_PATH # Where all projects will be contained

    SCRIPT_PATH = __file__
    SCRIPT_DIR_PATH = os.path.dirname(SCRIPT_PATH)
    PROJECTS_DB_PATH = os.path.join(SCRIPT_DIR_PATH, 'PROJECTS')

    if not os.path.exists(PROJECTS_DB_PATH):
        os.mkdir(PROJECTS_DB_PATH)
    #endregion

    # PROMPT USER TO CHOOSE PROJECT
    chosen_project = choose_project()

    #region # Intra-project paths
    # Project specific folder
    # lines folder
    # csv file
    global PROJECT_NAME
    global PROJECT_PATH
    global LINES_FOLDER_PATH
    global LINE_INFO_FILE_PATH
    # The existence of this file indicates a dir is a project
    global PROJECT_FLAG_FILE_PATH

    # Names and paths
    PROJECT_NAME = chosen_project
    PROJECT_PATH = os.path.join(PROJECTS_DB_PATH, PROJECT_NAME)
    PROJECT_FLAG_FILE_PATH = os.path.join(PROJECT_PATH, 'project.ACL')
    LINES_FOLDER_PATH = os.path.join(PROJECT_PATH, 'lines')
    LINE_INFO_FILE_PATH = os.path.join(PROJECT_PATH, 'line_data.csv')

    # Track info stuff
    global TRACKS_ARE_BGM
    TRACK_TYPE_HEADER = 'TRACK_TYPE:'
    TRACK_TYPE_OPTIONS = ['BGM', 'RADIO']

    # Check if paths exist
    if not os.path.exists(PROJECT_PATH):
        os.mkdir(PROJECT_PATH)

    if not os.path.exists(LINES_FOLDER_PATH):
        os.mkdir(LINES_FOLDER_PATH)

    if not os.path.exists(PROJECT_FLAG_FILE_PATH): # If project is new:
        TRACKS_ARE_BGM = get_track_playback_info(True) # Ask user track type
        with open(PROJECT_FLAG_FILE_PATH, 'w', encoding="utf-8") as file:
            # Write track type info:
            if TRACKS_ARE_BGM: # Tracks are BGM
                file.write(TRACK_TYPE_HEADER + TRACK_TYPE_OPTIONS[0])
            else: # Tracks are RADIO
                file.write(TRACK_TYPE_HEADER + TRACK_TYPE_OPTIONS[1])
    else: # If file already exists
        with open(PROJECT_FLAG_FILE_PATH, 'r', encoding="utf-8") as file:
            data = (file.read()).splitlines()
            # Find track type info
            track_type_entry = [x for x in data if TRACK_TYPE_HEADER in x][0]
            if TRACK_TYPE_OPTIONS[0] in track_type_entry:
                get_track_playback_info(False, True)
            else:
                get_track_playback_info(False, False)
        pass
    #endregion

def prompt_user_list(option_list, custom_answer):
    # This function creates a prompt to choose from a list.
    # Handles invalid answers. Answer must be an index.
    #
    # custom_answer: If true, allows for a custom answer
    # Returns index. If custom, return custom answer.

    if custom_answer:
        option_list.append('Other')

    for index, entry in enumerate(option_list):
        print('{} - {}'.format(index, entry))
    print()

    max_index = len(option_list)-1
    valid_index = True
    while True:
        answer = ''
        if not valid_index: # Display error message accordig to flag
            print('Input a valid index!')

        try:
            answer = int(input('Enter index: ')) # Receive answer
        except ValueError:
            valid_index = False # Set flag to false on invalidity
            continue

        if answer < 0 or answer > max_index: # Answer is outisde index range
            valid_index = False # Set flag to false on invalidity
            continue

        valid_index = True
        if custom_answer and answer == max_index:
            user_input = input('Input your answer: ')
            # Changing the list here changes the list passed as parameter.
            option_list.pop() # Remove old item
            option_list.append(user_input) # Insert custom answer
        break

    return answer

def choose_project():
    # Prompt user to choose which project to open/create
    # For now, name is hard-coded.
    global PROJECTS_DB_PATH

    projects_list = os.listdir(PROJECTS_DB_PATH)

    print('Choose which project to open/create:')

    chosen_project = projects_list[prompt_user_list(projects_list, True)]

    print('Opening "{}".'.format(chosen_project))

    return chosen_project

def get_rows_line_data_csv():
    # Opens csv and returns all rows
    with open(LINE_INFO_FILE_PATH, mode='r', encoding='UTF8', newline='') as line_file: # Open file
        csv_reader = csv.reader(line_file, delimiter='\\', quotechar='`')
        return [x for x in csv_reader]

def write_row_line_data_csv(data, is_list, append):
    # Write to the csv file.
    #
    # If append == True, append to file.
    # Else, overwrite it.
    #
    # If is_list, write all items of iterables as rows
    # Else, write a single row only

    mode = ''
    if append:
        mode = 'a'
    else:
        mode = 'w'

    with open(LINE_INFO_FILE_PATH, mode=mode, encoding='UTF8', newline='') as line_file: # Create file
            line_writer = csv.writer(line_file, delimiter='\\', quotechar='`')
            if is_list:
                line_writer.writerows(data)
            else:
                line_writer.writerow(data)

def handle__tracks_info_file():
    # Go through the lines in the index
    global PROJECT_PATH
    global LINES_FOLDER_PATH
    global LINE_INFO_FILE_PATH

    global INDEXING_CRITERIA

    global INPUT_EXIT_MESSAGE
    global INPUT_CONTINUE_MESSAGE

    track_list = os.listdir(LINES_FOLDER_PATH)

    # What info will be requested to index each line.
    # In the sublists, default values for each criteria, if they are a closed set.
    INDEXING_CRITERIA = ['CHARACTER',
                        'MISSION_NUMBER',
                        ['ACE_STYLE', ['MERCENARY', 'SOLDIER', 'KNIGHT', 'NONE']],
                        'TEXT']

    file_header = []
    # How many of the first rows shouldn't be editable
    constant_columns = 2
    # Init .csv file
    # If csv does not exist and there are tracks to be indexed:
    if (not os.path.exists(LINE_INFO_FILE_PATH)) and len(track_list) != 0:
        # Define data to be written to new file
        # Assemble header with criteria as columns
        for criteria in INDEXING_CRITERIA:
            if type(criteria) == str:
                file_header.append(criteria)
            else: # If criteria has preset answers:
                file_header.append(criteria[0])
        # Constant columns. "constant_columns" = 2, since there are 2
        file_header.insert(0, 'ID') # Insert ID column. It autoincrements.
        file_header.insert(1, 'FILENAME') # Insert FILENAME column.

        # Assemble all files
        all_rows = []
        for field_index, track in enumerate(track_list):
            row = [field_index, track]
            # Write empty data to criteria to have
            # all columns appear as items in the list when reading later.
            padding = ['' for x in INDEXING_CRITERIA]
            row += padding
            all_rows.append(row)

        # Write data to file
        write_row_line_data_csv(file_header, False, True) # Write header
        write_row_line_data_csv(all_rows, True, True)
    elif len(track_list) == 0: # If there are no tracks to be indexed:
        print('There are no tracks to be indexed!')
        print('Copy tracks to the "{}" folder'
              .format(os.path.basename(LINES_FOLDER_PATH)))
        input(INPUT_EXIT_MESSAGE)
        return False

    line_file_dump = []
    line_pure_info = []
    line_file_dump = get_rows_line_data_csv()

    # From here, pass "line_pure_info" as a parameter
    # to an indexing prompt function,
    # which will return an updated list.
    # The current function will take it and save it to the .csv file.
    edited_file_data = surf_lines(line_file_dump, True)

    # Rewrite file
    write_row_line_data_csv(edited_file_data, True, False)

def surf_lines(line_info, ignore_unknowns):
    # This will surf the lines and see what data needs completion.
    # Receives the "line_info" of read file, and value for "ignore_unknowns".
    #
    # "ignore_unknowns": During field completion,
    # skip values marked with '?', that couldn't be determined by the user.
    global INDEXING_CRITERIA

    global INPUT_CONTINUE_MESSAGE
    # Types of special data that can be found in a field
    empty_data = '' # Empty. Must be filled out.
    unknown_data = '?' # Seen, but not known yet. Must be skipped.
    dummy_data = 'dummy' # Marks dummy files. Skip & mark whole line as dummy.

    # Types of special data that can be input in a field
    entry_dummy = dummy_data # Same as before
    entry_skip_line = '||' # Skip whole line
    entry_unknown_line = '?:' # Mark whole line as unknown
    entry_terminate = '\\' # Char that terminates a prompt when entered

    num_static_columns = 2 # Ammount of static columns (ID, file name)

    #region # Special inputs
    # Print help screen that show what commands can be given:
    help_messages = [
        '\nTime to fill out the data!',
        'Special commands:',
        '"COMMAND" - Description'.format(),

        '{} - Dummy data. The whole line will be marked as such'.format(dummy_data),
        '{} - Skip line. Skip whole line without inputting anything'.format(entry_skip_line),
        '{} - Unknown data. When you don\'t know the answer for a field'.format(unknown_data),
        '{} - Unknown line. Mark whole line as unknown'.format(entry_unknown_line),
        '{} - Terminate. Stop inputting data, save and quit application'.format(entry_terminate),]
    for help in help_messages:
        print(help)

    input(INPUT_CONTINUE_MESSAGE)
    #endregion

    # Split data into parts
    header = line_info[0] # Get header
    line_info = line_info[1:] # Remove header from list.
    data_set = line_info # Pass info to 'data_set'.
    static_columns = data_set

    # Remove ID and track name from editable data
    data_set = [x[num_static_columns:] for x in data_set]
    # Reference header with only columns names for editable data
    editable_header = [x for x in header[num_static_columns:]]
    # Remove editable data from the 'static_columns'
    static_columns = [x[:num_static_columns] for x in static_columns]

    quit = False

    for line_index, line in enumerate(data_set):
        if quit: # If user chose to quit in previous line
            break
        playback = True
        track_name = static_columns[line_index][1]

        for field_index, field in enumerate(line): # Check field
            current_field = editable_header[field_index]

            # Decide what to do based on data
            if dummy_data in line: # Mark whole line and skip.
                data_set[line_index] = [dummy_data for x in line]
                break
            elif field == unknown_data: # Determine what to do
                if ignore_unknowns:
                    continue # Skip field
                else:
                    pass # Go ahead to prompt
            elif field == empty_data: # Must be filled up
                pass # Go ahead to prompt
            else: # Is user-input data. Skip.
                continue # Skip field

            if playback: # This is executed only once per track
                # Play track for user
                print('\n' + track_name) # Print file name

                track_path = os.path.join(LINES_FOLDER_PATH, track_name)
                play_track(track_path) # Play track
            playback = False

            # Data entry
            answer = input('Enter your data for {}: '.format(current_field))

            if answer == entry_terminate:
                # Quit prompt
                quit = True
                break
            elif answer == entry_dummy:
                # Mark whole line as dummy and skip
                data_set[line_index] = [dummy_data for x in line]
                break
            elif answer == entry_skip_line: # Skip line without inputting
                break
            elif answer == entry_unknown_line: # Mark whole line as unknown and skip
                data_set[line_index] = [unknown_data for x in line]
                break
            else: # Pass answer to field
                data_set[line_index][field_index] = answer
                pass

    # Reassemble file for outputting
    output_data = []
    for index, entry in enumerate(data_set):
        output_data.append(static_columns[index] + data_set[index])
    output_data.insert(0, header) # Insert header back

    return output_data

def get_track_playback_info(new_file, is_bgm = True):
    # This function prompts the use
    # If tracks are BGM, choose BGM parameters.
    # Else, use RADIO parameters.
    #
    # If new_file == true, prompt user to input type data
    # Else, receive if tracks are BGM or not

    # Argument list for MFAudio, with indexes for the lists used in this code
        # 0 -  /IFnnnnn	Input frequency
        # 1 -  /ICn	Input channels
        # 2 -  /IIxxxx	Input interleave (hex)
        # 3 -  /IHxxxx	Input headerskip (hex)
        # 4 -  /OTtttt	Output type (WAVU, VAGC,
        # 	            SS2U, SS2C, RAWU, RAWC)
        # 5 -  /OFnnnnn	Output frequency
        # 6 -  /OCn	Output channels
        # 7 -  /OIxxxx	Output interleave (hex)
        # 8 -  "InputFile"	Input file to play/convert
        # 9 -  "OutputFile"	Output file to convert to

    global TRACKS_ARE_BGM
    global MFAUDIO_ARG_SET
    global MFAUDIO_PATH
    MFAUDIO_PATH = os.path.join(os.path.dirname(SCRIPT_PATH), 'MFAudio', 'MFAudio.exe')

    # Argument set for BGM and RADIO tracks.
    # AFAIK, there is no reason for them to differ between AC5 and ACZ.
    MFAUDIO_BGM_ARGS = [44100, 2, 320, 0, 'WAVU', 44100, 2, 320]
    MFAUDIO_RADIO_ARGS = [22050, 1, 320, 0, 'WAVU', 22050, 1, 320]

    TYPE_OPTIONS = ['BGM', 'RADIO']

    if new_file: # If file is new, prompt user.
        print('Are you working with BGM or RADIO?')
        answer = prompt_user_list(TYPE_OPTIONS, False)
        if answer == 0:
            TRACKS_ARE_BGM = True
            MFAUDIO_ARG_SET = MFAUDIO_BGM_ARGS
        else:
            TRACKS_ARE_BGM = False
            MFAUDIO_ARG_SET = MFAUDIO_RADIO_ARGS
    else: # If file already exists, get type data from parameter
        if not is_bgm:
            TRACKS_ARE_BGM = False
            MFAUDIO_ARG_SET = MFAUDIO_RADIO_ARGS
        else:
            TRACKS_ARE_BGM = True
            MFAUDIO_ARG_SET = MFAUDIO_BGM_ARGS

def play_track(track_path):
    # This method plays the chosen track through MFAudio,
    # with the predefined arguments chosen on project creation
    global MFAUDIO_ARG_SET
    global SCRIPT_PATH
    global MFAUDIO_PATH

    input_filename = os.path.basename(track_path)

    args = ['/IF', '/IC', '/II', '/IH', '/OT', '/OF', '/OC', '/OI'] #Contains the core arguments.

    # Concatenate args and values
    args = [args[index] + str(MFAUDIO_ARG_SET[index]) for index, x in enumerate(args)]
    args.append('"{}"'.format(track_path)) # Append track path to args
    args.insert(0, MFAUDIO_PATH)

    ## Adds a input filename to 'argumentBuffer' string and...
    #argument_buffer = exe_filename + ' ' + argument_buffer + '"' + DESTINATION_PATH + '\\'+ input_filename + '"'
    #
    #if mode == 0:  # ... if the mode is "Convert (0)", add output filename as well.
    #    argument_buffer = argument_buffer + ' ' + '"' + output_filename + '"'
    #
    args_string = ''
    for index, arg in enumerate(args):
        args_string += arg
        if index < 9:
            args_string += ' '

    subprocess.Popen(args_string)

    pass

INPUT_EXIT_MESSAGE = 'PRESS ENTER TO EXIT'
INPUT_CONTINUE_MESSAGE = 'PRESS ENTER TO CONTINUE'

if __name__ == "__main__":
    main()


from time import sleep
from typing import List, Tuple
from pygame.surface import Surface
import pygame

from map import Map, Point

WHITE = (255, 255, 255)
BLUE = (0, 0, 255)
GREEN = (0, 255, 0)

def init_pygame(dimension: Tuple[int, int]) -> Surface:
  pygame.init()

  logo = pygame.image.load("../assets/logo32x32.png")
  pygame.display.set_icon(logo)
  pygame.display.set_caption("Drone Exploration with Evolutionary Algorithm")

  screen = pygame.display.set_mode(dimension)
  screen.fill(WHITE)
  return screen

def pygame_main_loop() -> None:
  while pygame.QUIT not in map(lambda e: e.type, pygame.event.get()):
    pygame.display.flip()

  pygame.quit()

def map_image(m: Map) -> Surface:
  image = Surface((20 * m.width, 20 * m.height))
  brick = Surface((20, 20))

  image.fill(WHITE)
  brick.fill(BLUE)

  for (x, y) in m.positions():
    if m[(x, y)] == 1:
      image.blit(brick, (y * 20, x * 20))

  return image

def show_map(m: Map) -> None:
  screen = init_pygame((m.width * 20, m.height * 20))
  screen.blit(map_image(m), (0, 0))
  pygame_main_loop()

def draw_path(m: Map, path: List[Point], speed: float = 1, markSeen: bool = True) -> None:
  screen = init_pygame((m.width * 20, m.height * 20))
  drona = pygame.image.load("../assets/drona.png")

  brick = Surface((20,20))
  brick.fill(GREEN)

  screen.blit(map_image(m), (0, 0))
  pygame.display.flip()
  sleep(0.5 * speed)

  for p in path:
    x, y = p

    if markSeen:
      for vx, vy in m.visible_area(p):
        screen.blit(brick, (vy * 20, vx * 20))

    screen.blit(drona, (y * 20, x * 20))
    pygame.display.flip()

    sleep(0.5 * speed)

  pygame_main_loop()


import urllib.request
import math

def Degree(WindDirComp):
    if 0 <= WindDirection < 5.62:
        WindDirComp = "N (North)"
    elif 5.63 < WindDirection < 16.87:
        WindDirComp = "NbE (North by East)"
    elif 16.88 < WindDirection < 28.12:
        WindDirComp = "NNE (North-Northeast)"
    elif 28.13 < WindDirection < 39.37:
        WindDirComp = "NEbN (Northeast by North)"
    elif 39.38 < WindDirection < 50.62:
        WindDirComp = "NE (Northeast)"
    elif 50.63 < WindDirection < 61.87:
        WindDirComp = "NEbE (Northeast by East)"
    elif 61.88 < WindDirection < 73.12:
        WindDirComp = "ENE (East-Northeast)"
    elif 73.13 < WindDirection < 84.37:
        WindDirComp = "EbN (East by North)"
    elif 84.38 < WindDirection < 95.62:
        WindDirComp = "E (East)"
    elif 95.63 < WindDirection < 106.87:
        WindDirComp = "EbS (East by South)"
    elif 106.88 < WindDirection < 118.12:
        WindDirComp = "ESE (East-Southeast)"
    elif 118.13 < WindDirection < 129.37:
        WindDirComp = "SEbE (Southeast by East)"
    elif 129.38 < WindDirection < 140.62:
        WindDirComp = "SE (Southeast)"
    elif 140.63 < WindDirection < 151.87:
        WindDirComp = "SEbS (Southeast by South)"
    elif 151.88 < WindDirection < 163.12:
        WindDirComp = "SSE (South-Southeast)"
    elif 163.13 < WindDirection < 174.37:
        WindDirComp = "SbE (South by East)"
    elif 174.38 < WindDirection < 185.62:
        WindDirComp = "S (South)"
    elif 185.63 < WindDirection < 196.87:
        WindDirComp = "SbW (South by West)"
    elif 196.88 < WindDirection < 208.12:
        WindDirComp = "SSW (South-Southwest)"
    elif 208.13 < WindDirection < 219.37:
        WindDirComp = "SWbS (Southwest by South)"
    elif 219.38 < WindDirection < 230.62:
        WindDirComp = "SW (Southwest)"
    elif 230.63 < WindDirection < 241.87:
        WindDirComp = "SWbW (Southwest by West)"
    elif 241.88 < WindDirection < 253.12:
        WindDirComp = "WSW (West-Southwest)"
    elif 253.13 < WindDirection < 264.37:
        WindDirComp = "WbS (West by South)"
    elif 264.38 < WindDirection < 275.62:
        WindDirComp = "W (West)"
    elif 275.63 < WindDirection < 286.87:
        WindDirComp = "WbN (West by North)"
    elif 286.88 < WindDirection < 298.12:
        WindDirComp = "WNW (West-Northwest)"
    elif 298.13 < WindDirection < 309.37:
        WindDirComp = "NWbW (Northwest by West)"
    elif 309.38 < WindDirection < 320.62:
        WindDirComp = "NW (Northwest)"
    elif 320.63 < WindDirection < 331.87:
        WindDirComp = "NWbN (Northwest by North)"
    elif 331.88 < WindDirection < 343.12:
        WindDirComp = "NNW (North-Northwest)"
    elif 343.13 < WindDirection < 354.37:
        WindDirComp = "NbW (North by West)"
    elif 354.38 < WindDirection <= 360:
        WindDirComp = "N (North)"
    return WindDirComp

ICAO = input("What airport would you like to find the weather conditions for? (Use ICAO code)")
ICAO = ICAO.upper()

pageText = urllib.request.urlopen('http://w1.weather.gov/data/METAR/' + ICAO + '.1.txt')
pageText = str(pageText.read())

A = pageText[45:]
print(A)
list = A.split()
Airport = list[0]

Day = list[1][0:2]
Hour = list[1][2:4]
Minute = list[1][4:6]

WindDir = list[2][0:3]
if WindDir == "VRB":
    WindDir = "Variable"
elif WindDir == "CLM":
    WindDir = "Calm"
else:
    WindDirDeg = WindDir + "°" #u'\u00b0'
    WindDirection = int(WindDir)
    WindDirComp = Degree(WindDir)

WindSpeed = int(list[2][3:5])
if list[2][5] == "K":
    WindMeasurement = "Knots"
    WindMPH = round(WindSpeed * 1.15078)
elif list[2][5] == "M":
    WindMeasurement = "Meters per Second"
    WindMPH = round(WindSpeed * 2.23694)

Visibility = list[3][:-2]

#Weather Types
if list[4][0:1] == "-":
    Weather1 = "Light"
elif list[4][0:1] == "+":
    Weather1 = "Heavy"
elif list[4][0:2] == "VC":
    Weather1 = "In the vicinity"
elif list[4][0:2] == "MI":
    Weather1 = "Shallow"
elif list[4][0:2] == "PR":
    Weather1 = "Partial"
elif list[4][0:2] == "BC":
    Weather1 = "Patches"
elif list[4][0:2] == "DR":
    Weather1 = "Low drifting"
elif list[4][0:2] == "BL":
    Weather1 = "Blowing"
elif list[4][0:2] == "SH":
    Weather1 = "Showers"
elif list[4][0:2] == "TS":
    Weather1 = "Thunderstorms"
elif list[4][0:2] == "FZ":
    Weather1 = "Freezing"
elif list[4][0:2] == "RA":
    Weather1 = "Rain"
elif list[4][0:2] == "DZ":
    Weather1 = "Drizzle"
elif list[4][0:2] == "SN":
    Weather1 = "Snow"
elif list[4][0:2] == "SG":
    Weather1 = "Snow Grains"
elif list[4][0:2] == "IC":
    Weather1 = "Ice Crystals"
elif list[4][0:2] == "PL":
    Weather1 = "Ice Pellets"
elif list[4][0:2] == "GR":
    Weather1 = "Hail"
elif list[4][0:2] == "GS":
    Weather1 = "Small Hail"
elif list[4][0:2] == "UP":
    Weather1 = "Unknown Precipitation"
elif list[4][0:2] == "FG":
    Weather1 = "Fog"
elif list[4][0:2] == "VA":
    Weather1 = "Volcanic Ash"
elif list[4][0:2] == "BR":
    Weather1 = "Broken Clouds"
elif list[4][0:2] == "HZ":
    Weather1 = "Haze"
elif list[4][0:2] == "DU":
    Weather1 = "Widespread Dust"
elif list[4][0:2] == "FU":
    Weather1 = "Smoke"
elif list[4][0:2] == "SA":
    Weather1 = "Sand"
elif list[4][0:2] == "PY":
    Weather1 = "Spray"
elif list[4][0:2] == "SQ":
    Weather1 = "Squalls"
elif list[4][0:2] == "PO":
    Weather1 = "Dust or Sand Whirls"
elif list[4][0:2] == "DS":
    Weather1 = "Duststorm"
elif list[4][0:2] == "SS":
    Weather1 = "Sandstorm"
elif list[4][0:2] == "FC":
    Weather1 = "Funnel Cloud"
elif list[4][0:3] == "CLR":
    Weather1 = "Clear"
else:
    Weather1 = "Unknown Weather"

print("Airport: " + Airport)
print("Day of Month is: " + Day + "th")
print("Time of Reading: " + Hour + ":" + Minute + " GMT")
print("Wind Direction: " + WindDirDeg)
print("                " + WindDirComp)
print("Wind Speed: " + str(WindSpeed) + " " + WindMeasurement)
print("            " + str(WindMPH) + " Miles per Hour")
print("Visibility: " + Visibility + " Miles")
print("The current weather is " + Weather1)

print("")
print("Press [Enter] to exit.")
input("> ")


#!/usr/bin/env python3
from simtk.openmm.app import *
from simtk.openmm import *
from simtk.unit import *
from sys import stdout
from pdbfixer import *
import mdtraj as md
from Bio.PDB.Polypeptide import *
from Bio.PDB.PDBParser import PDBParser
from Bio.PDB import PDBList
from Bio.PDB import PDBIO
from itertools import product, combinations
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
import os
import textwrap

from functionTerms.basicTerms import *
from functionTerms.contactTerms import *
from functionTerms.hydrogenBondTerms import *
from functionTerms.qBiasTerms import *
from functionTerms.debyeHuckelTerms import *
from functionTerms.templateTerms import *
from functionTerms.membraneTerms import *

se_map_3_letter = {'ALA': 0,  'PRO': 1,  'LYS': 2,  'ASN': 3,  'ARG': 4,
                   'PHE': 5,  'ASP': 6,  'GLN': 7,  'GLU': 8,  'GLY': 9,
                   'ILE': 10, 'HIS': 11, 'LEU': 12, 'CYS': 13, 'MET': 14,
                   'SER': 15, 'THR': 16, 'TYR': 17, 'VAL': 18, 'TRP': 19}

def identify_terminal_residues(pdb_filename):
    # identify terminal residues
    parser = PDBParser()
    structure = parser.get_structure('X', pdb_filename)
    terminal_residues = {}
    for model in structure:
        for chain in model:
            residues = list(chain.get_residues())
            terminal_residues[chain.id] = (residues[0].id[1], residues[-1].id[1])
        return terminal_residues

def prepare_pdb(pdb_filename, chains_to_simulate, use_cis_proline=False, keepIds=False, removeHeterogens=True):
    # for more information about PDB Fixer, see:
    # http://htmlpreview.github.io/?https://raw.github.com/pandegroup/pdbfixer/master/Manual.html
    # fix up input pdb
    cleaned_pdb_filename = "%s-cleaned.pdb" % pdb_filename[:-4]
    input_pdb_filename = "%s-openmmawsem.pdb" % pdb_filename[:-4]

    fixer = PDBFixer(filename=pdb_filename)

    # remove unwanted chains
    chains = list(fixer.topology.chains())
    chains_to_remove = [i for i, x in enumerate(chains) if x.id not in chains_to_simulate]
    fixer.removeChains(chains_to_remove)

    #Identify Missing Residues
    fixer.findMissingResidues()
    fixer.missingResidues = {}

    #Replace Nonstandard Residues
    fixer.findNonstandardResidues()
    fixer.replaceNonstandardResidues()

    #Remove Heterogens
    if removeHeterogens:
        fixer.removeHeterogens(keepWater=False)

    #Add Missing Heavy Atoms
    fixer.findMissingAtoms()
    fixer.addMissingAtoms()

    #Add Missing Hydrogens
    fixer.addMissingHydrogens(7.0)
    PDBFile.writeFile(fixer.topology, fixer.positions, open(cleaned_pdb_filename, 'w'), keepIds=keepIds)

    #Read sequence
    structure = PDBParser().get_structure('X', cleaned_pdb_filename)

    # identify terminal residues
    terminal_residues = identify_terminal_residues(cleaned_pdb_filename)

    # process pdb for input into OpenMM
    #Selects only atoms needed for the awsem topology
    output = open(input_pdb_filename, 'w')
    counter=0
    for line in open(cleaned_pdb_filename):
        # print(line)
        splitline = line.split()
        if len(line)>4 and line[0:4] == "ATOM":
            try:
                atom_index=line[6:11].strip()
                atom_type=line[12:16].strip()
                res_type=line[17:20].strip()
                chain=line[21].strip()
                res_index=line[22:26].strip()
                x=line[30:38].strip()
                y=line[38:46].strip()
                z=line[46:54].strip()
                element = line[76:78].strip()
            except ValueError:
                print(line)
                raise
        else:
            continue
        awsem_atoms = ["CA", "O", "CB", "C", "H", "N"]
        if int(res_index) == terminal_residues[chain][0]:
            awsem_atoms.remove("N")
            awsem_atoms.remove("H")
        if int(res_index) == terminal_residues[chain][1]:
            awsem_atoms.remove("C")
        # GLY should not has CB.
        if res_type == "GLY":
            awsem_atoms.remove("CB")
        if atom_type in awsem_atoms:
            line=list(line)
            if res_type == "GLY":
                line[17:20] = "IGL"
            elif res_type == "PRO":
                line[17:20] = "IPR"
            else:
                line[17:20] = "NGP"
            if atom_type == "CB":
                line[77] = "B"
            line=''.join(line)
            output.write(line)
            counter+=1
    #print("The system contains %i atoms"%counter)
    output.close()

    #Fix Virtual Site Coordinates:
    # prepare_virtual_sites(input_pdb_filename, use_cis_proline=use_cis_proline)
    prepare_virtual_sites_v2(input_pdb_filename, use_cis_proline=use_cis_proline)
    return input_pdb_filename, cleaned_pdb_filename

def prepare_virtual_sites_v2(pdb_file, use_cis_proline=False):
    parser = PDBParser(QUIET=True)
    structure=parser.get_structure('X',pdb_file,)
    res = list(structure.get_residues())
    model = structure[0]
    output_file = pdb_file
    f = open(pdb_file)
    all_lines = f.readlines()
    f.close()
    output = open(output_file, "w")
    index = 0
    for line in all_lines:
        # print(line)
        splitline = line.split()
        if len(line)>4 and line[0:4] == "ATOM":
            try:
                atom_index=line[6:11].strip()
                atom_type=line[12:16].strip()
                res_type=line[17:20].strip()
                chain=line[21].strip()
                res_index=line[22:26].strip()
                x=line[30:38].strip()
                y=line[38:46].strip()
                z=line[46:54].strip()
                element = line[76:78].strip()
            except ValueError:
                print(line)
                raise
        else:
            continue
        res_index = int(res_index)

        r_im = model[chain][max(res_index-1,1)]
        r_i = model[chain][res_index]
        try:
            r_ip = model[chain][res_index+1]
        except:
            r_ip = model[chain][res_index]  # won't be used
        if use_cis_proline and res_type == "IPR":
            n_coord = -0.2094*r_im['CA'].get_coord()+ 0.6908*r_i['CA'].get_coord() + 0.5190*r_im['O'].get_coord()
            c_coord = 0.2196*r_i['CA'].get_coord()+ 0.2300*r_ip['CA'].get_coord() + 0.5507*r_i['O'].get_coord()
            h_coord = -0.9871*r_im['CA'].get_coord()+ 0.9326*r_i['CA'].get_coord() + 1.0604*r_im['O'].get_coord()
        else:
            n_coord = 0.48318*r_im['CA'].get_coord()+ 0.70328*r_i['CA'].get_coord()- 0.18643 *r_im['O'].get_coord()
            c_coord = 0.44365*r_i['CA'].get_coord()+ 0.23520*r_ip['CA'].get_coord()+ 0.32115 *r_i['O'].get_coord()
            h_coord = 0.84100*r_im['CA'].get_coord()+ 0.89296*r_i['CA'].get_coord()- 0.73389 *r_im['O'].get_coord()

        line_list=list(line)
        index += 1
        line_list[6:11] = "{:5d}".format(index)
        if atom_type == "N":
            line_list[30:38] = '{:.8s}'.format('{:8.3f}'.format(n_coord[0]))
            line_list[38:46] = '{:.8s}'.format('{:8.3f}'.format(n_coord[1]))
            line_list[46:54] = list('{:.8s}'.format('{:8.3f}'.format(n_coord[2])))
        if atom_type == "C":
            line_list[30:38] = '{:.8s}'.format('{:8.3f}'.format(c_coord[0]))
            line_list[38:46] = '{:.8s}'.format('{:8.3f}'.format(c_coord[1]))
            line_list[46:54] = list('{:.8s}'.format('{:8.3f}'.format(c_coord[2])))
        if atom_type == "H":
            line_list[30:38] = '{:.8s}'.format('{:8.3f}'.format(h_coord[0]))
            line_list[38:46] = '{:.8s}'.format('{:8.3f}'.format(h_coord[1]))
            line_list[46:54] = list('{:.8s}'.format('{:8.3f}'.format(h_coord[2])))
        new_line=''.join(line_list)
        output.write(new_line)
    index += 1
    line_list[6:11] = "{:5d}".format(index)
    line_list[0:4] = "TER "
    line_list[30:78] = " " * 48
    new_line=''.join(line_list)
    output.write(new_line)
    output.write("END\n")
    output.close()

def prepare_virtual_sites(pdb_file, use_cis_proline=False):
    parser = PDBParser(QUIET=True)
    structure=parser.get_structure('X',pdb_file,)
    for model in structure:
        for chain in model:
            r_im={}
            r_i={}
            for residue in chain:
                r_im=r_i
                r_i={}
                for atom in residue:
                    r_i[atom.get_name()]=atom
                if use_cis_proline and residue.get_resname() == "IPR":
                    if 'N' in r_i:
                        r_i['N'].set_coord(-0.2094*r_im['CA'].get_coord()+ 0.6908*r_i['CA'].get_coord() + 0.5190*r_im['O'].get_coord())
                    if 'C' in r_im:
                        r_im['C'].set_coord(0.2196*r_im['CA'].get_coord()+ 0.2300*r_i['CA'].get_coord() + 0.5507*r_im['O'].get_coord())
                    if 'H' in r_i:
                        r_i['H'].set_coord(-0.9871*r_im['CA'].get_coord()+ 0.9326*r_i['CA'].get_coord() + 1.0604*r_im['O'].get_coord())
                else:
                    if 'N' in r_i:
                        r_i['N'].set_coord(0.48318*r_im['CA'].get_coord()+ 0.70328*r_i['CA'].get_coord()- 0.18643 *r_im['O'].get_coord())
                    if 'C' in r_im:
                        r_im['C'].set_coord(0.44365*r_im['CA'].get_coord()+ 0.23520*r_i['CA'].get_coord()+ 0.32115 *r_im['O'].get_coord())
                    if 'H' in r_i:
                        r_i['H'].set_coord(0.84100*r_im['CA'].get_coord()+ 0.89296*r_i['CA'].get_coord()- 0.73389 *r_im['O'].get_coord())
    io = PDBIO()
    io.set_structure(structure)
    io.save(pdb_file)

def build_lists_of_atoms(nres, residues):
    # build lists of atoms, residue types, and bonds
    atom_types=['n', 'h', 'ca', 'c', 'o', 'cb']
    res_types=[]
    atom_lists=dict(zip(atom_types,[[] for i in range(len(atom_types))]))
    for residue in residues:
        res_types.append(residue.name)
        atom_types=['n', 'h', 'ca', 'c', 'o', 'cb']
        residue_atoms = [x.index for x in residue.atoms()]
        if residue.index == 0:
            atom_types.remove('n')
            atom_lists['n'].append(-1)
            atom_types.remove('h')
            atom_lists['h'].append(-1)
        if residue.index == nres-1:
            atom_types.remove('c')
            atom_lists['c'].append(-1)
            pass
        if residue.name == "IGL" and 'cb' in atom_types:
            atom_types.remove('cb')
            atom_lists['cb'].append(-1)
        if residue.name in "IPR" and 'h' in atom_types:
            atom_types.remove('h')
            atom_lists['h'].append(-1)
        assert len(residue_atoms)==len(atom_types), '%s\n%s'%(str(residue_atoms),str(atom_types))
        atom_types=[a.name.lower() for a in residue._atoms] #Sometimes the atom order may be different
        for atom, atype in zip(residue_atoms, atom_types):
                atom_lists[atype].append(atom)

    return atom_lists, res_types

def build_lists_of_atoms_2(nres, residues, atoms):
    res_id = 0
    n = h = ca = c = o = cb = -1
    atom_types=['n', 'h', 'ca', 'c', 'o', 'cb']
    atom_types_table = {'N':'n', 'H':'h', 'CA':'ca', 'C':'c', 'O':'o', 'CB':'cb'}
    res_types=[]
    for residue in residues:
        res_types.append(residue.name)

    atom_lists=dict(zip(atom_types,[[-1]*nres for i in range(len(atom_types))]))
    for atom in atoms:
        atom_lists[atom_types_table[atom.name]][atom.residue.index] = atom.index

    return atom_lists, res_types

def ensure_atom_order(input_pdb_filename, quiet=1):
    # ensure order of ['n', 'h', 'ca', 'c', 'o', 'cb']
    # to be more specific, this ensure 'ca' always show up before 'c'.
    def first(t):
        return t[0]
    order_table = {'N':0, 'H':1, 'CA':2, 'C':3, 'O':4, 'CB':5}
    one_residue = []
    with open("tmp.pdb", "w") as out:
        with open(input_pdb_filename, "r") as f:
            all_lines = f.readlines()
            pre_res_id = all_lines[0].split()[5]
            for line in all_lines:
                info = line.split()
                if info[0]!="ATOM":
                    sorted_residue = sorted(one_residue, key=first)
                    for a in sorted_residue:
                        out.write(a[1])
                    one_residue = []
                    out.write(line)
                    continue
                res_id = info[5]
                if res_id != pre_res_id:
                    pre_res_id = res_id
                    # sort them in order
                    sorted_residue = sorted(one_residue, key=first)
                    for a in sorted_residue:
                        out.write(a[1])
                    if sorted_residue != one_residue:
                        if not quiet:
                            print("Reorder atom position")
                            print("Original, Changed to")
                        for i, t in zip(one_residue, sorted_residue):
                            if not quiet:
                                print(i[2], i[3], ", ", t[2], t[3])
                    one_residue = []
                atomType = info[2]
                one_residue.append((order_table[atomType], line, info[1], atomType))
            sorted_residue = sorted(one_residue, key=first)
            for a in sorted_residue:
                out.write(a[1])
    os.system(f"mv tmp.pdb {input_pdb_filename}")

def get_chain_starts_and_ends(all_res):
    chain_starts = []
    chain_ends = []
    pos = -1
    for i, res in enumerate(all_res):
        if res.chain.index != pos:
            chain_starts.append(i)
            pos = res.chain.index
            if i > 0:
                chain_ends.append(i-1)
    chain_ends.append(len(all_res)-1)
    return chain_starts, chain_ends

def setup_virtual_sites(nres, system, n, h, ca, c, o, cb, res_type, chain_starts, chain_ends, use_cis_proline=False):
    # set virtual sites
    for i in range(nres):
        if use_cis_proline and res_type[i] == "IPR":
            if i not in chain_starts:
                n_virtual_site = ThreeParticleAverageSite(ca[i-1], ca[i], o[i-1],
                                                         -0.2094, 0.6908, 0.5190)
                system.setVirtualSite(n[i], n_virtual_site)
                if not res_type[i] == "IPR":
                    h_virtual_site = ThreeParticleAverageSite(ca[i-1], ca[i], o[i-1],
                                                            -0.9871, 0.9326, 1.0604)
                    system.setVirtualSite(h[i], h_virtual_site)
            if i not in chain_ends:
                c_virtual_site = ThreeParticleAverageSite(ca[i], ca[i+1], o[i],
                                                        0.2196, 0.2300, 0.5507)
                # print("Virtual", c[i])
                system.setVirtualSite(c[i], c_virtual_site)
        else:
            if i not in chain_starts:
                n_virtual_site = ThreeParticleAverageSite(ca[i-1], ca[i], o[i-1],
                                                        0.48318, 0.70328, -0.18643)
                system.setVirtualSite(n[i], n_virtual_site)
                if not res_type[i] == "IPR":
                    h_virtual_site = ThreeParticleAverageSite(ca[i-1], ca[i], o[i-1],
                                                            0.84100, 0.89296, -0.73389)
                    system.setVirtualSite(h[i], h_virtual_site)
            if i not in chain_ends:
                c_virtual_site = ThreeParticleAverageSite(ca[i], ca[i+1], o[i],
                                                        0.44365, 0.23520, 0.32115)
                # print("Virtual", c[i])
                system.setVirtualSite(c[i], c_virtual_site)

def setup_bonds(nres, n, h, ca, c, o, cb, res_type, chain_starts, chain_ends):
    bonds = []
    for i in range(nres):
        bonds.append((ca[i], o[i]))
        if not res_type[i] == "IGL":
            bonds.append((ca[i], cb[i]))
        if  i not in chain_ends:
            bonds.append((ca[i], ca[i+1]))
            bonds.append((o[i], ca[i+1]))

    for i in range(nres):
        if i not in chain_starts and not res_type[i] == "IGL":
            bonds.append((n[i], cb[i]))
        if i not in chain_ends and not res_type[i] == "IGL":
            bonds.append((c[i], cb[i]))
        if i not in chain_starts and i not in chain_ends:
            bonds.append((n[i], c[i]))
    return bonds

def read_fasta(fastaFile):
    with open(fastaFile) as input_data:
        data = ""
        for line in input_data:
            if(line[0] == ">"):
                print(line)
            elif(line == "\n"):
                pass
            else:
                data += line.strip("\n")
    return data

def formatResidue_ThreeLetterCodeToOne(residue):
    ThreeToOne = {'ALA':'A','ARG':'R','ASN':'N','ASP':'D','CYS':'C','GLU':'E','GLN':'Q','GLY':'G','HIS':'H',
        'ILE':'I','LEU':'L','LYS':'K','MET':'M','PHE':'F','PRO':'P','SER':'S','THR':'T','TRP':'W',
        'TYR':'Y','VAL':'V'}
    residue_name = residue.get_resname()
    try:
        oneLetter = ThreeToOne[residue_name]
    except:
        print(f"Unknown residue: {residue.get_full_id()}, treat as ALA")
        residue_name = "ALA"
    return ThreeToOne[residue_name]

def getSeqFromCleanPdb(input_pdb_filename, chains='A', writeFastaFile=False):
    cleaned_pdb_filename = input_pdb_filename.replace("openmmawsem.pdb", "cleaned.pdb")
    pdb = input_pdb_filename.replace("-openmmawsem.pdb", "")
    fastaFile = pdb + ".fasta"

    s = PDBParser().get_structure("X", cleaned_pdb_filename)
    m = s[0] # model 0
    seq = ""
    if writeFastaFile:
        with open(fastaFile, "w") as out:
            for chain in chains:
                out.write(f">{pdb.upper()}:{chain}\n")
                c = m[chain]
                chain_seq = ""
                for residue in c:
                    chain_seq += formatResidue_ThreeLetterCodeToOne(residue)
                out.write("\n".join(textwrap.wrap(chain_seq, width=80))+"\n")
                seq += chain_seq
    else:
        for chain in chains:
            c = m[chain]
            chain_seq = ""
            for residue in c:
                chain_seq += formatResidue_ThreeLetterCodeToOne(residue)
            seq += chain_seq
    return seq

def getSeq(input_pdb_filename, chains='A', writeFastaFile=False, fromPdb=False, fromFasta=None):
    if fromPdb:
        cleaned_pdb_filename = input_pdb_filename.replace("openmmawsem.pdb", "cleaned.pdb")
        pdb = input_pdb_filename.replace("-openmmawsem.pdb", "")
        fastaFile = pdb + ".fasta"

        s = PDBParser().get_structure("X", cleaned_pdb_filename)
        m = s[0] # model 0
        seq = ""
        if writeFastaFile:
            with open(fastaFile, "w") as out:
                for chain in chains:
                    out.write(f">{pdb.upper()}:{chain.upper()}\n")
                    c = m[chain]
                    chain_seq = ""
                    for residue in c:
                        chain_seq += formatResidue_ThreeLetterCodeToOne(residue)
                    out.write("\n".join(textwrap.wrap(chain_seq, width=80))+"\n")
                    seq += chain_seq
        else:
            for chain in chains:
                c = m[chain]
                chain_seq = ""
                for residue in c:
                    chain_seq += formatResidue_ThreeLetterCodeToOne(residue)
                seq += chain_seq
    elif fromFasta:
        seq = ""
        with open(fromFasta) as f:
            for line in f:
                if line[0] == ">":
                    pass
                else:
                    # print(line)
                    seq += line.strip()
    return seq

def fixPymolPdb(location):
    # location = "1r69.pdb"
    with open("tmp", "w") as out:
        with open(location, "r") as f:
            for line in f:
                info = list(line)
                if len(info) > 21:
                    info[21] = "A"
                out.write("".join(info))
    os.system(f"mv tmp {location}")

def download(pdb_id):
    if not os.path.isfile(f"{pdb_id}.pdb"):
        PDBList().retrieve_pdb_file(pdb_id.lower(), pdir='.', file_format='pdb')
        os.rename("pdb%s.ent" % pdb_id, f"{pdb_id}.pdb")

class OpenMMAWSEMSystem:
    def __init__(self, pdb_filename, chains='A', xml_filename='awsem.xml', k_awsem=1.0, seqFromPdb=None, includeLigands=False):
        # read PDB
        self.pdb = PDBFile(pdb_filename)
        self.forcefield = ForceField(xml_filename)
        if not includeLigands:
            self.system = self.forcefield.createSystem(self.pdb.topology)
            # define convenience variables
            self.nres = self.pdb.topology.getNumResidues()
            self.natoms = self.pdb.topology.getNumAtoms()
            self.residues = list(self.pdb.topology.residues())
            self.resi = [x.residue.index for x in list(self.pdb.topology.atoms())]
            # build lists of atoms and residue types
            # self.atom_lists,self.res_type=build_lists_of_atoms(self.nres, self.residues)
            self.atom_lists,self.res_type=build_lists_of_atoms_2(self.nres, self.residues, self.pdb.topology.atoms())
        if includeLigands:
            print(self.pdb.topology)
            [templates, names] = self.forcefield.generateTemplatesForUnmatchedResidues(self.pdb.topology)
            # a = templates[0]
            for a in templates:
                for a1 in a.atoms:
                    # we can define the types of ligand atoms using their names.
                    a1.type = a1.name
                    # if a1.element.symbol == "C":
                    #     a1.type = "CA"
                    # else:
                    #     a1.type = a1.element.symbol
                    # a1.type = a1.element.symbol
                self.forcefield.registerResidueTemplate(a)

            res_list = list(self.pdb.topology.residues())
            atom_list = list(self.pdb.topology.atoms())
            protein_resNames = ["NGP", "IGL", "IPR"]
            DNA_resNames = ["DA", "DC", "DT", "DG"]
            protein_res_list = []
            DNA_res_list = []
            ligand_res_list = []
            for res in res_list:
                if res.name in protein_resNames:
                    protein_res_list.append(res)
                elif res.name in DNA_resNames:
                    DNA_res_list.append(res)
                else:
                    ligand_res_list.append(res)

            protein_atom_list = []
            DNA_atom_list = []
            ligand_atom_list = []
            for atom in atom_list:
                if atom.residue.name in protein_resNames:
                    protein_atom_list.append(atom)
                elif atom.residue.name in DNA_resNames:
                    DNA_atom_list.append(atom)
                else:
                    ligand_atom_list.append(atom)
            self.ligand_res_list = ligand_res_list
            self.ligand_atom_list = ligand_atom_list
            self.protein_atom_list = protein_atom_list
            self.system = self.forcefield.createSystem(self.pdb.topology)
            self.nres = len(protein_res_list)
            self.residues = protein_res_list
            print(f"Number of residues: {self.nres}, Number of ligands: {len(ligand_res_list)}")
            # self.natoms = len(protein_atom_list)
            self.natoms = self.pdb.topology.getNumAtoms()
            self.resi = [x.residue.index if x in protein_atom_list else -1 for x in atom_list]
            self.atom_lists,self.res_type=build_lists_of_atoms_2(self.nres, self.residues, protein_atom_list)

        # print(self.atom_lists,self.res_type)
        self.n =self.atom_lists['n']
        self.h =self.atom_lists['h']
        self.ca=self.atom_lists['ca']
        self.c =self.atom_lists['c']
        self.o =self.atom_lists['o']
        self.cb=self.atom_lists['cb']

        self.chain_starts, self.chain_ends = get_chain_starts_and_ends(self.residues)
        # setup virtual sites
        # if Segmentation fault in setup_virtual_sites, use ensure_atom_order
        setup_virtual_sites(self.nres, self.system, self.n, self.h, self.ca, self.c, self.o,
                            self.cb, self.res_type, self.chain_starts, self.chain_ends)
        # setup bonds
        self.bonds = setup_bonds(self.nres, self.n, self.h, self.ca, self.c, self.o,
                            self.cb, self.res_type, self.chain_starts, self.chain_ends)
        # identify terminal_residues
        # self.terminal_residues = identify_terminal_residues(pdb_filename)
        # set overall scaling
        self.k_awsem = k_awsem
        # keep track of force names for output purposes
        self.force_names = []
        # save seq info
        if seqFromPdb is None:
            self.seq = getSeq(pdb_filename, chains=chains, fromPdb=True)
        else:
            self.seq = seqFromPdb
        # elif seqFromPdb == 0:
        #     self.seq = getSeq(pdb_filename, chains=chains, fromFasta=True)

    def addForces(self, forces):
        for i, (force) in enumerate(forces):
            self.addForce(force)
            force.setForceGroup(i+1)

    def addForce(self, force):
        self.system.addForce(force)

    def addForcesWithDefaultForceGroup(self, forces):
        for i, (force) in enumerate(forces):
            self.addForce(force)

def read_trajectory_pdb_positions(pdb_trajectory_filename):
    import uuid, os
    pdb_trajectory_contents = open(pdb_trajectory_filename).read().split("MODEL")[1:]
    pdb_trajectory_contents = ['\n'.join(x.split('\n')[1:]) for x in pdb_trajectory_contents]
    pdb_trajectory = []
    for i, pdb_contents in enumerate(pdb_trajectory_contents):
        temporary_file_name = str(uuid.uuid4())
        temporary_pdb_file = open(temporary_file_name, 'w')
        temporary_pdb_file.write(pdb_contents)
        temporary_pdb_file.close()
        pdb = PDBFile(temporary_file_name)
        pdb_trajectory.append(pdb)
        os.remove(temporary_file_name)
    return pdb_trajectory

def compute_order_parameters(openmm_awsem_pdb_file, pdb_trajectory_filename, order_parameters,
            platform_name='CPU', k_awsem=1.0, compute_mdtraj=False, rmsd_reference_structure=None,
            compute_total_energy=True, energy_columns=None, xml_filename="awsem.xml"):
    pdb_trajectory = read_trajectory_pdb_positions(pdb_trajectory_filename)
    order_parameter_values = []
    for i, order_parameter in enumerate(order_parameters):
        order_parameter_values.append([])
        oa = OpenMMAWSEMSystem(openmm_awsem_pdb_file, k_awsem=k_awsem, xml_filename=xml_filename)
        platform = Platform.getPlatformByName(platform_name) # OpenCL, CUDA, CPU, or Reference
        integrator = VerletIntegrator(2*femtoseconds)
        oa.addForce(order_parameter)
        simulation = Simulation(oa.pdb.topology, oa.system, integrator, platform)
        for pdb in pdb_trajectory:
            simulation.context.setPositions(pdb.positions)
            state = simulation.context.getState(getEnergy=True)
            order_parameter_values[i].append(state.getPotentialEnergy().value_in_unit(kilojoule_per_mole))
    if compute_mdtraj:
        md_traj_order_parameters = compute_mdtraj_order_parmeters(pdb_trajectory_filename, rmsd_reference_structure=rmsd_reference_structure)
        for key, value in md_traj_order_parameters.items():
            if len(value.flatten()) == len(pdb_trajectory):
                order_parameter_values.append(value)
    if compute_total_energy:
        order_parameter_values.append([])
        for i in range(len(pdb_trajectory)):
            order_parameter_values[-1].append(np.sum([order_parameter_values[x-1][i] for x in energy_columns]))
    if not compute_mdtraj:
        return np.array(order_parameter_values)
    else:
        return np.array(order_parameter_values), md_traj_order_parameters

def compute_mdtraj_order_parmeters(trajectory_file, rmsd_reference_structure=None):
    # documentation: http://mdtraj.org/1.8.0/analysis.html#
    trajectory = md.load(trajectory_file)

    return_values = []
    return_value_names = []

    if not rmsd_reference_structure == None:
        reference = md.load(rmsd_reference_structure)
        rmsd = md.rmsd(trajectory, reference)
        return_values.append(rmsd)
        return_value_names.append("RMSD")

    hydrogen_bonds = np.array([np.sum(x) for x in md.kabsch_sander(trajectory)])
    return_values.append(hydrogen_bonds)
    return_value_names.append("HBondEnergy")

    ss = md.compute_dssp(trajectory)
    shape = ss.shape
    transdict = dict(zip(list(set(list(ss.flatten()))),range(len(list(set(list(ss.flatten())))))))
    ss = np.array([transdict[x] for x in ss.flatten()]).reshape(shape).T
    return_values.append(ss)
    return_value_names.append("SecondaryStructure")

    rg = md.compute_rg(trajectory)
    return_values.append(rg)
    return_value_names.append("Rg")

    distances, residue_pairs = md.compute_contacts(trajectory, scheme='ca')
    contacts = md.geometry.squareform(distances, residue_pairs)
    return_values.append(contacts)
    return_value_names.append("Contacts")

    return dict(zip(return_value_names, return_values))

def plot_free_energy(order_parameter_file, labels, bins=20, discard=.2, twodlimits=[[0,1],[0,1]]):
    import matplotlib as mpl
    import scipy.ndimage
    mpl.rcParams['font.size'] = 24

    if len(labels) > 2:
        print("Too many labels to plot.")
        return
    order_parameters = np.loadtxt(order_parameter_file).T
    data_labels = open(order_parameter_file).readlines()[0].split()[1:]
    data = []
    for label in labels:
        raw_data = order_parameters[data_labels.index(label)]
        data_after_discard = raw_data[int(len(raw_data)*discard):]
        comparison_data_set_size = (1.0-discard)/2
        data_set_1 = raw_data[int(len(raw_data)*discard):int(len(raw_data)*(discard+comparison_data_set_size))]
        data_set_2 = raw_data[int(len(raw_data)*(discard+comparison_data_set_size)):]
        data.append(data_after_discard)
    if len(labels) == 1:
        hist, bins = np.histogram(data[0], density=True, bins=bins)
        hist /= np.sum(hist)
        f = -np.log(hist)
        bin_centers = (bins[:-1]+bins[1:])/2
        hist, bins = np.histogram(data_set_1, density=True, bins=bins)
        hist /= np.sum(hist)
        f_1 = -np.log(hist)
        bin_centers_1 = (bins[:-1]+bins[1:])/2
        hist, bins = np.histogram(data_set_2, density=True, bins=bins)
        hist /= np.sum(hist)
        f_2 = -np.log(hist)
        bin_centers_2 = (bins[:-1]+bins[1:])/2
        plt.figure(figsize=(12,8))
        plt.plot(bin_centers, f, label="All data")
        plt.plot(bin_centers_1, f_1, label="Split data 1")
        plt.plot(bin_centers_2, f_2, label="Split data 2")
        plt.legend(bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.)
        plt.title(labels[0])
        plt.xlabel(labels[0])
        plt.ylabel("Free Energy (kT)")
    if len(labels) == 2:
        H, xedges, yedges = np.histogram2d(data[0], data[1], bins=bins)
        xcenters = (xedges[:-1] + xedges[1:])/2.0
        ycenters = (yedges[:-1] + yedges[1:])/2.0
        H /= np.sum(H)
        H = -np.log(H)
        H -= np.min(H)

        fig = plt.figure(figsize=(12, 10))
        ax = fig.add_subplot(111)
        ax.set_title("%s vs. %s" % (labels[0], labels[1]))
        ax.set_xlim(twodlimits[0])
        ax.set_ylim(twodlimits[1])
        ax.set_xlabel(labels[0])
        ax.set_ylabel(labels[1])
        X, Y = np.meshgrid(xcenters, ycenters)
        plt.contourf(X, Y, H, cmap=plt.cm.get_cmap('jet'))

        plt.colorbar()
        plt.tight_layout()
        plt.show()

def compute_perturbed_energies(openmm_awsem_pdb_file, pdb_trajectory_filename, perturbations, order_parameter_values, platform_name='CPU', k_awsem=1.0, total_energy_column=15):
    pdb_trajectory = read_trajectory_pdb_positions(pdb_trajectory_filename)
    all_perturbed_energies = []
    for i, perturbation in enumerate(perturbations):
        # compute new energies
        perturbed_energies = []
        for j, energy_term in enumerate(perturbation):
            perturbed_energies.append([])
            oa = OpenMMAWSEMSystem(openmm_awsem_pdb_file, k_awsem=k_awsem)
            platform = Platform.getPlatformByName(platform_name) # OpenCL, CUDA, CPU, or Reference
            integrator = VerletIntegrator(2*femtoseconds)
            oa.addForce(energy_term[1])
            simulation = Simulation(oa.pdb.topology, oa.system, integrator, platform)
            for pdb in pdb_trajectory:
                simulation.context.setPositions(pdb.positions)
                state = simulation.context.getState(getEnergy=True)
                perturbed_energies[j].append(state.getPotentialEnergy().value_in_unit(kilojoule_per_mole))
        perturbed_energies = np.array(perturbed_energies)
        total_perturbed_energies = np.sum(perturbed_energies, axis=0)
        # compute new total energy by subtracting original value and adding new value
        new_total_energy = np.array(order_parameter_values[:, total_energy_column-1])
        for j, energy_term in enumerate(perturbation):
            new_total_energy -= order_parameter_values[:, energy_term[0]-1]
        new_total_energy += total_perturbed_energies
        all_perturbed_energies.append(new_total_energy)
    all_perturbed_energies = np.array(all_perturbed_energies)
    return all_perturbed_energies.T

def pick_structures(label, conditions_string, metadata_file, reference_structure, num_snapshots=6000, order_parameter_file_name="order_parameters.txt", pdb_trajectory_filename="output.pdb"):
    # parse restrictions (including which file to pull columns from)
    # read in data (including extra files if necessary)
    # go through the data and filter out snapshots that do not satisfy the criteria
    # select a subset of the structures that satisfy the constraints

    def parse_conditions_string(conditions_string):
        conditions = []
        condition_signs = []
        conditions_string = conditions_string.split()
        for condition in conditions_string:
            if "gt" in condition:
                condition_signs.append("+")
                condition = condition.split("gt")
            if "lt" in condition:
                condition_signs.append("-")
                condition = condition.split("lt")
            conditions.append(condition)
        return conditions, condition_signs

    def load_all_data(metadata_file, conditions):
        data_files = list(np.loadtxt(metadata_file, dtype=str)[:,0])
        num_files = len(data_files)
        num_conditions = len(conditions)
        # Load all data into array
        data_array = np.zeros((num_files, num_conditions, num_snapshots))
        for i, data_file in enumerate(data_files):
            all_order_parameters = np.loadtxt(data_file)
            for j, condition in enumerate(conditions):
                data_array[i][j] = all_order_parameters[0:num_snapshots,int(condition[0])-1]

        data_array = np.swapaxes(data_array,1,2)
        return data_array

    def write_selected_pdbs(pdb_files, snapshots_in_pdb_files):
        structure_index = 1
        selected_pdbs = open("%s.pdb" % label, 'w')
        for pdb_file in pdb_files:
            pdb_trajectory_contents = open(pdb_file).read().split("MODEL")[1:]
            pdb_trajectory_contents = ['\n'.join(x.split('\n')[1:]) for x in pdb_trajectory_contents]
            for snapshot in snapshots_in_pdb_files[pdb_files.index(pdb_file)]:
                selected_pdbs.write("MODEL        %d\n" % structure_index)
                selected_pdbs.write(pdb_trajectory_contents[snapshot])
                structure_index += 1
        selected_pdbs.close()

    # Lists
    files_array = list(np.loadtxt(metadata_file, dtype=str)[:,0])
    conditions, condition_signs = parse_conditions_string(conditions_string)
    data_array = load_all_data(metadata_file, conditions)

    # File names and parameters
    output_file_name = label + ".dat"

    structure_index = 1

    # loop over data and output those points that satisfy all conditions
    output_file = open(output_file_name, "w")
    pdb_files = []
    snapshots_in_pdb_files = []

    # loop over files
    for i, data_file in enumerate(files_array):
        # loop over snapshots
        for j, snapshot in enumerate(data_array[i]):
            bad_condition = False
            # loop over conditions
            for k, condition in enumerate(conditions):
                condition_boundary = float(condition[1])
                # If all conditions are satisfied, print out the data
                if condition_signs[k] == "+":
                    if not data_array[i][j][k] > condition_boundary: bad_condition = True
                elif condition_signs[k] == "-":
                    if not data_array[i][j][k] < condition_boundary: bad_condition = True
                else:
                    print("Bad condition argument.")
                    sys.exit()
            if not bad_condition:
                output_file.write("%d %s %s\n" % (structure_index, data_file, j+1))
                pdb_file = data_file.replace(order_parameter_file_name, pdb_trajectory_filename)
                if not pdb_file in pdb_files:
                    pdb_files.append(pdb_file)
                    snapshots_in_pdb_files.append([])
                snapshots_in_pdb_files[pdb_files.index(pdb_file)].append(j)
                structure_index += 1
    output_file.close()
    write_selected_pdbs(pdb_files, snapshots_in_pdb_files)

    pymol_script = open("%s.pml" % label, 'w')
    pymol_script.write("load %s\n" % reference_structure)
    pymol_script.write("load_traj %s.pdb\n" % label)
    object_name = os.path.basename(reference_structure)[:-4]
    pymol_script.write("intra_fit %s\n" % object_name)
    pymol_script.write("smooth\n")
    pymol_script.close()


# 5)Write a program that asks the user to enter an integer and prints two integers, root and pwr, such that
# 0 < pwr < 6 and root**pwr is equal to the integer entered by the user. If no such pair of integers exists
# , it should print a message to that effect.
n=int(input())

for i in range(1,6):
	root=n**(1/i)
	print('For ',i)

	if int(root)**i==n:
		print ("root=",root,"pow=",i)
	else:
		print('not found')


from flask import Flask, render_template, request
import wordmatics as wm

app = Flask(__name__)

app.template_folder = '../templates'
app.static_folder='../static'

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/submit', methods=['POST'])
def submit():
    # Call the C++ executable with 'entry' as input
    # Process the result and return it
    letters = request.form.get('a')
    letters += request.form.get('b')
    letters += request.form.get('c')
    letters += request.form.get('d')
    letters += request.form.get('e')
    letters += request.form.get('f')
    center = request.form.get('i')
    s_words_list = ''
    f_words_list = ''
    #SLOW SLOW SLOW
    all =  wm.all_words_list()
    val = wm.valid_words_list()
    for word in all:
        f = wm.filter(word, letters, center)
        if f:
            s_words_list += f'
{word}
'
    for word in val:
        if wm.pangram(word, letters, center) and wm.filter(word, letters, center):
            f_words_list +=  f'
{word}
'
        else:
            if wm.filter(word, letters, center):
                f_words_list +=  f'
{word}
'
    #result is good

    return render_template("index.html", letters="Letters: " + letters, center="Center: " + center, valid_english_list=f_words_list, all_words_list=s_words_list)

if __name__ == '__main__':
    app.run(port=4174, debug=True)