import os
import shutil
import sys
import zipfile
from contextlib import suppress

if sys.version_info[:2] < (3, 8):
    sys.exit("You are running an incompatible version of Python, please use >= 3.8")

import argparse
import logging
import tarfile
from io import BytesIO

import urllib3

sys.path.append(os.path.join(os.path.abspath(os.path.dirname(__file__)), ".."))

import lib.cuckoo.common.colors as colors
from lib.cuckoo.common.constants import CUCKOO_ROOT
from lib.cuckoo.common.path_utils import path_exists, path_mkdir

blocklist = {}
if path_exists(os.path.join(CUCKOO_ROOT, "utils", "community_blocklist.py")):
    from utils.community_blocklist import blocklist

log = logging.getLogger(__name__)
URL = "https://github.com/kevoreilly/community/archive/{0}.tar.gz"

def flare_capa(proxy=None):
    signature_urls = (
        "https://github.com/mandiant/capa/raw/master/sigs/1_flare_msvc_rtf_32_64.sig",
        "https://github.com/mandiant/capa/raw/master/sigs/2_flare_msvc_atlmfc_32_64.sig",
        "https://github.com/mandiant/capa/raw/master/sigs/3_flare_common_libs.sig",
    )
    try:
        if proxy:
            http = urllib3.ProxyManager(proxy)
        else:
            http = urllib3.PoolManager()
        data = http.request("GET", "https://github.com/mandiant/capa-rules/archive/master.zip").data
        dest_folder = os.path.join(CUCKOO_ROOT, "data")
        shutil.rmtree((os.path.join(dest_folder, "capa-rules-master")), ignore_errors=True)
        shutil.rmtree((os.path.join(dest_folder, "capa-rules")), ignore_errors=True)
        zipfile.ZipFile(BytesIO(data)).extractall(path=dest_folder)
        os.rename(os.path.join(dest_folder, "capa-rules-master"), os.path.join(dest_folder, "capa-rules"))

        # shutil.rmtree((os.path.join(dest_folder, "capa-signatures")), ignore_errors=True)
        capa_sigs_path = os.path.join(dest_folder, "capa-signatures")
        if not os.path.isdir(capa_sigs_path):
            path_mkdir(capa_sigs_path)
        for url in signature_urls:
            signature_name = url.rsplit("/", 1)[-1]
            with http.request("GET", url, preload_content=False) as sig, open(
                os.path.join(capa_sigs_path, signature_name), "wb"
            ) as out_sig:
                shutil.copyfileobj(sig, out_sig)

        print("[+] FLARE CAPA rules/signatures installed")
    except Exception as e:
        print(e)

def mitre():
    """Urls might change, for proper urls see https://github.com/swimlane/pyattck"""
    try:
        from pyattck import Attck
    except ImportError:
        print("Missed dependency: install pyattck library, see requirements for proper version")
        return

    mitre = Attck(
        nested_techniques=True,
        use_config=False,
        save_config=False,
        config_file_path=os.path.join(CUCKOO_ROOT, "data", "mitre", "config.yml"),
        data_path=os.path.join(CUCKOO_ROOT, "data", "mitre"),
        enterprise_attck_json="https://raw.githubusercontent.com/mitre/cti/master/enterprise-attack/enterprise-attack.json",
        pre_attck_json="https://raw.githubusercontent.com/mitre/cti/master/pre-attack/pre-attack.json",
        mobile_attck_json="https://raw.githubusercontent.com/mitre/cti/master/mobile-attack/mobile-attack.json",
        ics_attck_json="https://raw.githubusercontent.com/mitre/cti/master/ics-attack/ics-attack.json",
        nist_controls_json="https://raw.githubusercontent.com/center-for-threat-informed-defense/attack-control-framework-mappings/master/frameworks/ATT%26CK-v9.0/nist800-53-r4/stix/nist800-53-r4-controls.json",
        generated_nist_json="https://swimlane-pyattck.s3.us-west-2.amazonaws.com/attck_to_nist_controls.json",
    )

    print("[+] Updating MITRE datasets")
    mitre.update()

def install(enabled, force, rewrite, filepath, access_token=None, proxy=False):
    if filepath and path_exists(filepath):
        t = tarfile.TarFile.open(filepath, mode="r:gz")
    else:
        print(f"Downloading modules from {URL}")
        try:
            if proxy:
                http = urllib3.ProxyManager(proxy)
            else:
                http = urllib3.PoolManager()
            if access_token is None:
                data = http.request("GET", URL).data
            elif "github" in URL:
                data = http.request(
                    "GET", URL, headers={"Authorization": f"token {access_token}", "User-Agent": "CAPEv2_sandbox"}
                ).data
            else:
                data = http.request("GET", URL, headers={"PRIVATE-TOKEN": access_token}).data

            if b"Not Found" == data:
                print("You don't have permissions to access this repo")
                sys.exit(-1)
            t = tarfile.TarFile.open(fileobj=BytesIO(data), mode="r:gz")
        except Exception as e:
            print("ERROR: Unable to download archive: %s" % e)
            sys.exit(-1)

    folders = {
        "feeds": "modules/feeds",
        "signatures": "modules/signatures",
        "processing": "modules/processing",
        "reporting": "modules/reporting",
        "machinery": "modules/machinery",
        "analyzer": "analyzer",
        "data": "data",
        "integrations": "lib/cuckoo/common/integrations",
    }

    members = t.getmembers()
    directory = members[0].name.split("/", 1)[0]

    for category in enabled:
        folder = folders.get(category, False)
        if not folder:
            continue

        print(f"\nInstalling {colors.cyan(category.upper())}")

        # E.g., "community-master/modules/signatures".
        name_start = f"{directory}/{folder}"
        for member in members:
            if not member.name.startswith(name_start) or name_start == member.name:
                continue

            filepath = os.path.join(CUCKOO_ROOT, folder, member.name[len(name_start) + 1 :])
            if member.name.lower().endswith((".gitignore", "readme.md", "-ci.yml")):
                continue

            if member.isdir():
                if not path_exists(filepath):
                    path_mkdir(filepath)
                continue

            if not rewrite and path_exists(filepath):
                print(f'File "{filepath}" already exists, {colors.yellow("skipped")}')
                continue

            install = False
            dest_file = os.path.basename(filepath)

            if filepath in blocklist.get(category, []):
                print(f'You have blacklisted file: {dest_file}. {colors.yellow("skipped")}')
                continue

            if not force:
                while True:
                    choice = input(f'Do you want to install file "{dest_file}"? [yes/no] ')
                    if choice.lower() in ("y", "yes"):
                        install = True
                        break
                    elif choice.lower() in ("n", "no"):
                        break
            else:
                install = True

            if install:
                if not path_exists(os.path.dirname(filepath)):
                    path_mkdir(os.path.dirname(filepath))

                print(f'File "{filepath}" {colors.green("installed")}')
                open(filepath, "wb").write(t.extractfile(member).read())

def main():
    global URL

    parser = argparse.ArgumentParser()
    parser.add_argument("-a", "--all", help="Download everything", action="store_true", required=False)
    parser.add_argument("-e", "--feeds", help="Download CAPE feed modules", action="store_true", required=False)
    parser.add_argument("-s", "--signatures", help="Download CAPE signatures", action="store_true", required=False)
    parser.add_argument("-p", "--processing", help="Download processing modules", action="store_true", required=False)
    parser.add_argument("-m", "--machinery", help="Download machine managers", action="store_true", required=False)
    parser.add_argument("-r", "--reporting", help="Download reporting modules", action="store_true", required=False)
    parser.add_argument("-an", "--analyzer", help="Download analyzer modules/binaries/etc", action="store_true", required=False)
    parser.add_argument("-data", "--data", help="Download data items", action="store_true", required=False)
    parser.add_argument("-i", "--integrations", help="Download integration items", action="store_true", required=False)
    parser.add_argument(
        "-f", "--force", help="Install files without confirmation", action="store_true", default=False, required=False
    )
    parser.add_argument("-w", "--rewrite", help="Rewrite existing files", action="store_true", required=False)
    parser.add_argument("-b", "--branch", help="Specify a different branch", action="store", default="master", required=False)
    parser.add_argument(
        "--file", help="Specify a local copy of a community .zip file", action="store", default=False, required=False
    )
    parser.add_argument(
        "-cr", "--capa-rules", help="Download capa rules and signatures", action="store_true", default=False, required=False
    )
    parser.add_argument("--mitre", help="Download updated MITRE JSONS", action="store_true", default=False, required=False)
    parser.add_argument(
        "-u", "--url", help="Download community modules from the specified url", action="store", default=None, required=False
    )
    parser.add_argument(
        "-t", "--token", help="Access token to download private repositories", action="store", default=None, required=False
    )
    parser.add_argument("--proxy", help="Proxy to use. Ex http://127.0.0.1:8080", action="store", required=False)
    args = parser.parse_args()

    URL = args.url or URL.format(args.branch)
    enabled = []

    if args.all:
        enabled = ["feeds", "processing", "signatures", "reporting", "machinery", "analyzer", "data", "integrations"]
        flare_capa()
    else:
        if args.feeds:
            enabled.append("feeds")
        if args.signatures:
            enabled.append("signatures")
        if args.processing:
            enabled.append("processing")
        if args.reporting:
            enabled.append("reporting")
        if args.machinery:
            enabled.append("machinery")
        if args.analyzer:
            enabled.append("analyzer")
        if args.data:
            enabled.append("data")
        if args.integrations:
            enabled.append("integrations")

    if args.capa_rules:
        flare_capa(args.proxy)
        if not enabled:
            return

    if args.mitre:
        mitre()
        if not enabled:
            return

    if not enabled:
        print(colors.red("You need to enable a category!\n"))
        parser.print_help()
        return

    install(enabled, args.force, args.rewrite, args.file, args.token, args.proxy)

if __name__ == "__main__":
    with suppress(KeyboardInterrupt):
        main()


#!/usr/bin/env python
# coding: utf-8
import numpy as np
import matplotlib.pylab as plt

Tolerance=0.
size=200
positive_charges=8
negative_charges=7

x_pos=np.random.randint(.05*size,.95*size,positive_charges)
y_pos=np.random.randint(.05*size,.95*size,positive_charges)

x_neg=np.random.randint(.05*size,.95*size,negative_charges)
y_neg=np.random.randint(.05*size,.95*size,negative_charges)

def geometry(mat):
    mat[:,0]=0
    mat[:,-1]=0
    mat[0,:]=0
    mat[-1,:]=0
    mat[x_pos,y_pos]=1
    mat[x_neg,y_neg]=-1

def time_step(mat):
    copia=np.zeros(mat.shape)
    copia[1:-1,1:-1]=(mat[2:,1:-1]+mat[0:-2,1:-1]+mat[1:-1,2:]+mat[1:-1,:-2])/4.
    geometry(copia)
    diff=(np.sqrt((mat-copia)**2).sum()/(100*100))
    return diff,copia

mat=2*np.random.random((size,size))-1
for t in range(1000000):
    diff,mat=time_step(mat)
    if t%5000==0:
        print(t,diff)
    if diff<=Tolerance:
        print(t,diff)
        break

pl=plt.imshow(mat,cmap='RdBu',vmin=-1,vmax=1)
plt.colorbar(pl)
plt.show()


import pygame
import numpy as np

width = 1200 #//2
height = 800 #//2
max_iteracion = 20
factorColor = 255//max_iteracion
dw = dh = 20

def iteraMandelbrot(x0,y0):
    global max_iteracion
    x = 0
    y = 0

    iteracion = 0

    while x*x + y*y < 4 and iteracion < max_iteracion:
        xtemp = x*x - y*y + x0
        y = 2*x*y + y0

        x = xtemp

        iteracion += 1
    return iteracion

def getX0Y0(i,j):
    x0 = (float(i)/width)*3.5 - 2.1
    y0 = (float(j)/height)*2.4 - 1.2
    return x0,y0

def getMandel(i,j):
    global mandel
    if mandel[i,j] == 0:
        x0, y0 = getX0Y0(i,j)
        iteracion = iteraMandelbrot(x0,y0)
        mandel[i,j] = iteracion
    else:
        iteracion = mandel[i,j]

    return iteracion

def addRect(x,y,dx,dy):
    global pendientes
    if not (x,y,dx,dy) in pendientes:
        pendientes.append((x,y,dx,dy))

def drawSquare(i,j,dx,dy):
    global fractal, factorColor
    if dx >= 2 and dy >= 2:
        iteracion00   = getMandel(i,   j)
        iteraciondx0  = getMandel(i+dx,j)
        iteracion0dy  = getMandel(i,   j+dy)
        iteraciondxdy = getMandel(i+dx,j+dy)

        if iteracion00 == iteraciondxdy and  iteracion00 == iteracion0dy and  iteracion00 == iteraciondx0 and dx <20:
            color = iteracion00 * factorColor
            if dx>50:
                print(f'{iteracion00} {dx} {dy} ')
        else:
            dx2 = dx//2
            dy2 = dy//2
            if dx2>0 and dy2>0:
                if iteracion00 != iteraciondx0:
                    addRect(i,     j,     dx2,dy2)
                    addRect(i+dx2, j,     dx2,dy2)
                if iteraciondx0 != iteraciondxdy:
                    addRect(i+dx2, j,     dx2,dy2)
                    addRect(i+dx2, j+dy2, dx2,dy2)
                if iteracion00 != iteracion0dy:
                    addRect(i,     j,     dx2,dy2)
                    addRect(i,     j+dy2, dx2,dy2)
                if iteracion0dy != iteraciondxdy:
                    addRect(i,     j+dy2, dx2,dy2)
                    addRect(i+dx2, j+dy2, dx2,dy2)
            else:
                print(f'{dx2} {dy2}')
            color = (iteracion00 + iteraciondxdy + iteracion0dy + iteraciondx0) * factorColor // 4
    else:
        iteracion = getMandel(i,j)
        color =iteracion * factorColor
    pygame.draw.rect(fractal,(color,color,color),(i,j,dx+1,dy+1))

pygame.init()

screen = pygame.display.set_mode((width,height),pygame.DOUBLEBUF | pygame.HWSURFACE)
pygame.display.set_caption("Mandelbrot Fractal")

fractal = screen.copy()

fractal.fill((0,0,0))

mandel = np.zeros((width + dw, height + dh))

pendientes = []
pendientesOld = []

for i in range(0,width,dw):
    for j in range(0,height,dh):
        drawSquare(i,j,dw,dh)

screen.blit(fractal, (0,0))
pygame.display.flip()

contador = 0
while len(pendientes) > 0:

    pendientesOld = pendientes
    pendientes = []
    # print(f'pendientes {len(pendientesOld)}')
    while len(pendientesOld)>0:
         x,y,dx,dy = pendientesOld.pop()

         contador += 1
         drawSquare(x,y,dx,dy)
         if contador%100 == 0:
             screen.blit(fractal, (0,0))
             pygame.display.flip()
             print(f'pendientes: {len(pendientesOld)}')

print(f'hemos terminado {len(pendientes)} {len(pendientesOld)}')
screen.blit(fractal, (0,0))
pygame.display.flip()
running = True

while running:

    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_ESCAPE:
                running = False

pygame.quit()


import numpy as npf

"""
You invest 1000 for 5 years an a yearly return of 8%
    @FutureValueOfAnInvestment
hence :
        Year 1 = 1000 + 8%
        Year 2 = Year 1 + 8%
        Year 3 = Year 2 + 8%
        Year 4 = Year 3 + 8%
        Year 5 = Year 4 + 8%
"""
#Calculate how much you'll have after 5 years

Years = int(input("Please Enter the number of years\n"))

Am = int(input("Please Enter the amount\n"))
Amount = Am * -1

Per = float(input("Please Enter the number the percentage e.g. (0.05)\n"))

PeriodicPayments = float(input("Please Enter the amount\n"))

result = npf.fv(rate = Per, nper = Years, pmt = 0, pv = Amount)

print(result)

"""

The Code -
    nper = number of periods/Year
    pv = present value (investment amount is subtracted since we are the ones that add it in)
    pmt =   periodic payments corresponds to investments/payments in our case its zero
            payment against loan principle plus interest.

"""


import regression as reg
import data as d
import numpy as np
import matplotlib.pyplot as plt
import random

df = d.read_file_columns()

train_df,valid_df,test_df = d.split_groups(df)

x_train_features,y_train_values = d.get_x_y(train_df)
x_valid_features,y_valid_values = d.get_x_y(valid_df)
x_test_features,y_test_values = d.get_x_y(test_df)

font1 = {'family':'serif','size':20}
font2 = {'family':'serif','size':15}

def algorithm1():
    epochs = 100
    num_tests = 6
    alpha_nums = []
    for i in range(num_tests):
        alpha = random.uniform(0.01,0.05)
        alpha_nums.append(alpha)
        error_train_list,error_valid_list = reg.regression(
            x_train_features,y_train_values,x_valid_features,y_valid_values,alpha,epochs)

        plt.figure()
        p1 = plt.scatter(x=list(range(epochs)), y=error_train_list,color="lightskyblue")
        p2 = plt.scatter(x=list(range(epochs)), y=error_valid_list,color="lightgreen")

        plt.legend((p1,p2),["Train error","Validation error"])

        plt.title("Train and Validation Error vs. Epoch",fontdict=font1)
        plt.xlabel("Epoch",fontdict=font2)
        plt.ylabel("Error",fontdict=font2)
        #plt.show()
        plt.savefig(f"plots/plot_{i}.png")

    textfile = open("alpha_nums.txt", "w")
    for element in alpha_nums:
        textfile.write(str(element) + "\n")
    textfile.close()

def algorithm2():
    epochs = 100
    alpha = 0.05
    error_train_list,error_valid_list,error_test_list, = reg.regression2(
        x_train_features,y_train_values,x_valid_features,y_valid_values,x_test_features,y_test_values,alpha,epochs)

    plt.figure()
    p1 = plt.scatter(x=list(range(epochs)), y=error_train_list,color="blue",alpha=0.75)
    p2 = plt.scatter(x=list(range(epochs)), y=error_valid_list,color="green",alpha=0.75)
    p3 = plt.scatter(x=list(range(epochs)), y=error_test_list,color="red",alpha=0.75)

    plt.legend((p1,p2,p3),["Train error","Validation error","Test error"])

    plt.title("Train, Validation and Test Error vs. Epoch",fontdict=font1)
    plt.xlabel("Epoch",fontdict=font2)
    plt.ylabel("Error",fontdict=font2)
    #plt.show()
    plt.savefig("plots/plot_test.png")

def algorithm3():
    epochs = 100
    alphas = [0.005, 0.5]
    for i in range(2):
        alpha = alphas[i]
        error_train_list,error_valid_list = reg.regression(
            x_train_features,y_train_values,x_valid_features,y_valid_values,alpha,epochs)

        plt.figure()
        p1 = plt.scatter(x=list(range(epochs)), y=error_train_list,color="lightskyblue")
        p2 = plt.scatter(x=list(range(epochs)), y=error_valid_list,color="lightgreen")

        plt.legend((p1,p2),["Train error","Validation error"])

        plt.title("Train and Validation Error vs. Epoch",fontdict=font1)
        plt.xlabel("Epoch",fontdict=font2)
        plt.ylabel("Error",fontdict=font2)
        #plt.show()
        plt.savefig(f"plots/plots_{i}.png")

#algorithm1()
#algorithm2()
#algorithm3()


# Kim Lang #33791771
# Eric Chou #95408627
# Jianing Tong #50419612

# CS 121 Project #2
import logging
import re
from urllib.parse import urlparse, urljoin
from corpus import Corpus
import os
from lxml import html
from frontier import Frontier

logger = logging.getLogger(__name__)

class Crawler:
    """
    This class is responsible for scraping urls from the next available link in frontier and adding the scraped links to
    the frontier
    """

    def __init__(self, frontier):
        self.frontier = frontier
        self.corpus = Corpus()
        self.store_dynamicURL = dict()  # Dict of dynamic URLS as key, freq as value
        self.store_URLs = dict()  # Dict of all subdomains visited as key, query as value
        self.outputURL = ("", 0)    #contains the url with the most output links and the count of output links
        self.downloadedURLs = []    #list of downloaded URLs
        self.traps = set()          #set of all traps to prevent showing repeat traps

    def start_crawling(self):
        """
        This method starts the crawling process which is scraping urls from the next available link in frontier and adding
        the scraped links to the frontier
        """
        while self.frontier.has_next_url():
            url = self.frontier.get_next_url()
            self.downloadedURLs.append(url)
            logger.info("Fetching URL %s ... Fetched: %s, Queue size: %s", url, self.frontier.fetched,
                        len(self.frontier))
            url_data = self.fetch_url(url)

            for next_link in self.extract_next_links(url_data):
                if self.corpus.get_file_name(next_link) is not None:
                    if self.is_valid(next_link):                #if the link is valid it stores into the frontier
                        self.frontier.add_url(next_link)
                    elif self.is_valid(next_link) == False:     #if the link is not valid it stores into the trap set()
                        self.traps.add(next_link)               #the way to identify a trap
        for url in self.frontier.urls_queue:                    #loops through the urls in queue
            self.downloadedURLs.append(url)                     #puts the urls that are downloaded into our downloadedURLS list
        self.write_analytics("analytics.txt")                   #writes the analytics to the analytics.txt file

    #writes the analytics required into our analytics.txt file
    def write_analytics(self, filename):
        with open(filename, 'w') as f:                  #opens the file up
            f.write("#1: SUBDOMAINS" + "\n")            #writes all the subdomains from our store_URLs
            for key, value in self.store_URLs.items():  #and writes the number of links from the subdomains
                f.write(key + "\t" + str(len(value)) + '\n')
            f.write("\n" + "#2 URL: " + self.outputURL[0] + " has the most output links of " + str(#writes the url with the most output links and the #
                self.outputURL[1]) + "\n")
            f.write("#3: DOWNLOADED LINKS" + "\n")
            for url in self.downloadedURLs:             #loops through the downloaded links and writes into the file
                f.write(url + "\n")
            f.write("\n" + "TRAPS:" + "\n")
            for trap in self.traps:                     #loops through the traps and writes the traps into the file
                f.write(trap + "\n")
        f.close()

    def fetch_url(self, url):
        """
        This method, using the given url, should find the corresponding file in the corpus and return a dictionary
        containing the url, content of the file in binary format and the content size in bytes
        :param url: the url to be fetched
        :return: a dictionary containing the url, content and the size of the content. If the url does not
        exist in the corpus, a dictionary with content set to None and size set to 0 can be returned.
        """
        #base url data
        url_data = {
            "url": url,
            "content": None,
            "size": 0
        }
        filepath = self.corpus.get_file_name(url)  # gets looks up for a local file and sends the address
        if filepath == None:    # if there is no filepath, return base url_data
            return url_data
        url_data["size"] = os.path.getsize(filepath) # gets the size of file using the filepath
        with open(filepath, "rb") as f:               # opens up the file and reads the file into binary format
            data = f.read()                             # reading
            url_data["content"] = data                 #places all the contents of the file in binary format into the content of the dict
        f.close()                                     #closes the file after done reading
        if "?" in url:                               #checks if the url in the dictionary is a dynamic url
            dynamic_url = url.split("?")[0]          #splits the url and indexing the first part before the ?
            dynamic_freq = self.store_dynamicURL.get(dynamic_url, 0)        #gets the frequency of the current dynamic url
            self.store_dynamicURL[dynamic_url] = dynamic_freq + 1           #adds one to the current frequency of the dynamic_url

        # Keeps track of subdomain visited for Analytics
        url_split = url.split("/")                                          #splits the url
        if url_split[2] in self.store_URLs.keys():                          #url_split[2] == the subdomain part
            self.store_URLs[url_split[2]].add("/".join(url_split[3:]))      # adds the path into the value of store_URLs
        else:
            self.store_URLs[url_split[2]] = set()
            self.store_URLs[url_split[2]].add("/".join(url_split[3:]))

        return url_data                                                   #returns the dict with the url, content in binary form, and size of the file

    def extract_next_links(self, url_data):
        """
        The url_data coming from the fetch_url method will be given as a parameter to this method. url_data contains the
        fetched url, the url content in binary format, and the size of the content in bytes. This method should return a
        list of urls in their absolute form (some links in the content are relative and needs to be converted to the
        absolute form). Validation of links is done later via is_valid method. It is not required to remove duplicates
        that have already been fetched. The frontier takes care of that.

        Suggested library: lxml
        """
        outputLinks = []
        tree = html.fromstring(url_data["content"])             #html.fromstring() expects the content in bytes  tree contains the html file
        url_list = tree.xpath('//a/@href')                      #xpath() scrapes the html file and gives us all the links we want in a list
        base_domain = url_data["url"]                           #the fetched_url
        for sub_url in url_list:                                                #goes through the urllist
            if sub_url.startswith("http://") or sub_url.startswith("https://"): #if the link is an absolute form add into outputlinks
                outputLinks.append(str(sub_url))
            else:                                                               #if not in absolute form or anything else
                joined_url = urljoin(base_domain, sub_url)                      #first join url with the fetched url then add into outputlinks
                outputLinks.append(str(joined_url))                             #joins the query or path

        if self.outputURL[1] < len(url_list):                                   # adds the next link if it has more links than the one before
            self.outputURL = (url_data["url"], len(url_list))

        return outputLinks

    def is_valid(self, url):
        """
        Function returns True or False based on whether the url has to be fetched or not. This is a great place to
        filter out crawler traps. Duplicated urls will be taken care of by frontier. You don't need to check for duplication
        in this method
        """
        parsed = urlparse(url)
        url_breakdown = url.split('/')  # list of url split by '/'

        # Checks for URLs with repetitions in the query and
        # returns false if greater than 2.
        url_dict = dict()
        for word in url_breakdown:
            freq = url_dict.get(word, 0)
            url_dict[word] = freq + 1
        for value in url_dict.values():
            if value > 2:               #if there is more than 2 duplicates in the url, it is considered as a trap
                return False

        # Checks for URLs that are super long
        if len(url_breakdown[-1]) > 150: # checking if the path is extremely long
            return False

        # Checks for dynamic URLs that impede progress of crawler
        dynamic_url = url.split("?")[0]
        if dynamic_url in self.store_dynamicURL.keys() and self.store_dynamicURL[dynamic_url] > 20:  #checks if the dynamic url appears more than 20 times
            return False

        if parsed.scheme not in set(["http", "https"]):
            return False

        try:
            return ".ics.uci.edu" in parsed.hostname \
                   and not re.match(".*\.(css|js|bmp|gif|jpe?g|ico" + "|png|tiff?|mid|mp2|mp3|mp4" \
                                    + "|wav|avi|mov|mpeg|ram|m4v|mkv|ogg|ogv|pdf" \
                                    + "|ps|eps|tex|ppt|pptx|doc|docx|xls|xlsx|names|data|dat|exe|bz2|tar|msi|bin|7z|psd|dmg|iso|epub|dll|cnf|tgz|sha1" \
                                    + "|thmx|mso|arff|rtf|jar|csv" \
                                    + "|rm|smil|wmv|swf|wma|zip|rar|gz|pdf)$", parsed.path.lower())

        except TypeError:
            print("TypeError for ", parsed)
            return False


# 이거 다시 봐야함;;
import heapq

node_length = 3
edge_length = 3

graph = [
    [],
    [(1, 2), (3, 3)],
    [(2, 3), (1, 1)],
    [(3, 1), (2, 2)],
]

visited = [False for _ in range(node_length + 1)]

def prim(start_node):
    q = [(0, start_node)]
    total_cost = 0

    while q:
        cost, node = heapq.heappop(q)

        if not visited[node]:
            visited[node] = True
            total_cost += cost

            for cost, edge in graph[node]:
                heapq.heappush(q, (cost, edge))

    return total_cost

print(visited)
print(prim(1))

.archive/2023-상반기-알고리즘-1/그래프/2. 최소 신장 트리(프림).py

2. 최소 신장 트리(프림).py


"""
The No-K-Means algorithm is a one-pass incremental clustering algorithm.
Breaking that down, it means that the algorithm:

1. Needs to read vectors only once to assign them to a cluster, and
2. Adds vectors to clusters, but does not remove them.

The algorithm compares incoming vectors to existing clusters.
The technique adds the vector to the most similar cluster if the similarity exceeds a certain threshold.
Otherwise, it creates a new cluster with the vector.

Therefore the complexity of the algorithm is :math:`O(nm)`, where :math:`n` is the number of vectors and :math:`m` is the number of clusters.
Since incremental approaches are often used in online settings, where timeliness is just as important as performance, the algorithm also has a freeze period.

The algorithm retires clusters that have not been updated after seeing a number of vectors, equivalent to the freeze period.
That means that if the freeze period is :math:`f`, and the cluster does not receive any vectors from any of the last :math:`f` vectors, the algorithm freezes it.
Frozen clusters cannot receive new vectors, and thus the algorithm can skip many checks, improving the efficiency.

.. note::

    This implementation is based on the algorithm outlined in `this paper by Azzopardi et al. in 2016 <https://www.researchgate.net/profile/Colin_Layfield/publication/303893387_Extended_No-K-Means_for_Search_Results_Clustering/links/575acd4208ae9a9c95518dfd.pdf>`_.
    The concept of one-pass incremental algorithms is very general, however, so you will find approaches like this one elsewhere.
"""

import os
import sys

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

from algorithms import ClusteringAlgorithm
from cluster import Cluster

class NoKMeans(ClusteringAlgorithm):
    """
    The No-K-Means algorithm class maintains two variables in its state:

    1. The clusters that are active, or which have not been frozen yet; and
    2. The frozen clusters, which may not receive new vectors.

    .. warning::

        By default, the algorithm does not store frozen clusters because they hog memory over time.
        You can specify that the algorithm should store frozen clusters by initializing the algorithm with the ``store_frozen`` parameter set to ``True``.
        Normally, algorithms (such as the :class:`~queues.consumers.eld_consumer.ELDConsumer`) can retain the clusters that they need themselves, and the algorithm removes the unused, frozen clusters.

    :ivar frozen_clusters: A list of clusters that have been retired.
    :vartype frozen_clusters: list of :class:`~vsm.clustering.cluster.Cluster` instances
    :ivar threshold: The minimum similarity between a vector and a cluster's centroid.
                     If any cluster exceeds this threshold, the vector is added to that cluster.
    :vartype threshold: float
    :ivar freeze_period: The number of vectors that arrive without being added to a cluster before it is frozen.
    :vartype freeze_period: int
    :ivar store_frozen: A boolean indicating whether frozen clusters should be retained.
    :vartype store_frozen: bool
    """

    def __init__(self, threshold, freeze_period, store_frozen=False):
        """
        Initialize the algorithm with an empty set of active and frozen clusters.
        The No-K-Means takes in two algorithm-related parameters:

        1. The threshold, or the minimum similarity to add a vector to a cluster; and
        2. The freeze period, or how long a cluster can go without receiving new vectors before being frozen.

        By default, the algorithm does not retain clusters once they are frozen, but you can change this behavior by setting the ``store_frozen`` parameter to ``True``.

        :param threshold: The minimum similarity between a vector and a cluster's centroid.
                          If any cluster exceeds this threshold, the vector is added to that cluster.
        :type threshold: float
        :param freeze_period: The number of vectors that arrive without being added to a cluster before it is frozen.
        :type freeze_period: int
        :param store_frozen: A boolean indicating whether frozen clusters should be retained.
        :type store_frozen: bool
        """

        super(NoKMeans, self).__init__()
        self.frozen_clusters = [ ]

        self.threshold = threshold
        self.freeze_period = freeze_period
        self.store_frozen = store_frozen

    def cluster(self, vectors, *args, **kwargs):
        """
        Cluster the given vectors.
        This function uses the ``clusters`` property as a list of active clusters.
        It automatically freezes inactive clusters as it iterates through the vectors.

        Any additional arguments and keyword arguments are passed on to the :func:`~vsm.cluster.cluster.Cluster.similarity` function.

        .. note::

            This function only returns the clusters that received new vectors during the function call.
            To get a list of _all_ active clusters, use the ``clusters`` property.

        :param vectors: The list of vectors to cluster.
        :type vectors: list of :class:`~vsm.vector.Vector`

        :return: The clusters that received vectors during this function call.
        :rtpye: list of :class:`~vsm.clustering.cluster.Cluster` instances
        """

        updated_clusters = [ ]

        for vector in vectors:
            """
            Freeze inactive clusters first.
            In this way, nothing gets added to them, thereby resetting their age.
            """
            for cluster in self.clusters:
                self._update_age(cluster)
                if self._to_freeze(cluster):
                    self._freeze(cluster)

            """
            If there are active clusters, get the closest cluster.
            If the vector's similarity with the cluster exceeds the threshold, add the vector to the cluster.
            The cluster's age is resetted.
            """
            if self.clusters:
                cluster, similarity = self._closest_cluster(vector, *args, **kwargs)
                if similarity >= self.threshold:
                    cluster.vectors.append(vector)
                    self._reset_age(cluster)
                    updated_clusters.append(cluster)
                    continue

            """
            If there was no similar cluster, create a new cluster with just that vector.
            """
            cluster = Cluster([ vector ])
            self.clusters.append(cluster)
            updated_clusters.append(cluster)

        return list(set(updated_clusters))

    def _update_age(self, cluster, increment=1):
        """
        Update the age of the given cluster.

        :param cluster: The cluster whose age will be incremented.
        :type cluster: :class:`~vsm.clustering.cluster.Cluster`
        :param increment: The amount by which to increment the age.
        :type increment: int
        """

        cluster.attributes['age'] = cluster.attributes.get('age', 0) + increment

    def _to_freeze(self, cluster):
        """
        Check if the cluster's age is high enough that the cluster should be frozen.
        A cluster should be frozen if its age is longer than the freeze period.

        :param cluster: The cluster whose age will be incremented.
        :type cluster: :class:`~vsm.clustering.cluster.Cluster`

        :return: A boolean indicating whether the cluster should be frozen.
        :rtype: bool
        """

        return cluster.attributes.get('age') > self.freeze_period

    def _freeze(self, cluster):
        """
        Freeze the given cluster.
        This removes the cluster from the list of active clusters.
        Simultaneously, it is added to the list of frozen clusters.

        :param cluster: The cluster whose age will be incremented.
        :type cluster: :class:`~vsm.clustering.cluster.Cluster`

        :raises ValueError: When the cluster is not active.
        """

        self.clusters.remove(cluster)

        if self.store_frozen:
            self.frozen_clusters.append(cluster)

    def _closest_cluster(self, vector, *args, **kwargs):
        """
        Get the closest cluster to the vector and return the similarity score.
        This function only compares the vector with active clusters.
        Any additional arguments and keyword arguments are passed on to the :func:`~vsm.cluster.cluster.Cluster.similarity` function.

        :param vector: The vector to compare with all clusters.
        :type vector: :class:`~vsm.vector.Vector`

        :return: A tuple consisting of the closest cluster and the similarity between it and the vector.
                 If there are no active clusters, ``None`` is returned instead.
        :rtype: tuple or None
        """

        if not self.clusters:
            return None

        similarities = [ (cluster, cluster.similarity(vector, *args, **kwargs)) for cluster in self.clusters ]
        closest_cluster = sorted(similarities, key=lambda tuple: tuple[1], reverse=True)[0]
        return closest_cluster

    def _reset_age(self, cluster):
        """
        Reset the cluster age to zero.
        This function is called when an existing cluster receives a new vector.

        :param cluster: The cluster whose age will be incremented.
        :type cluster: :class:`~vsm.clustering.cluster.Cluster`
        """

        cluster.attributes['age'] = 0


import orca
import numpy as np
import pandas as pd

# data documentation: https://berkeley.app.box.com/notes/282712547032

# Set data directory

d = '/home/data/spring_2019/2025/'

if 'data_directory' in orca.list_injectables():
    d = orca.get_injectable('data_directory')

# b = '/home/data/spring_2019/base/'
############################################################

# Tables from MTC Bay Area UrbanSim

@orca.table(cache=True)
def parcels():
    try:
        df = pd.read_csv(
            d + 'parcels.csv',
            index_col='primary_id',
            dtype={'primary_id': int, 'block_id': str})
    except:
        df = pd.read_csv(
            d + 'parcels.csv',
            index_col='parcel_id',
            dtype={'parcel_id': int, 'block_id': str})
    return df

@orca.table(cache=True)
def buildings():
    df = pd.read_csv(
        d + 'buildings.csv',
        index_col='building_id', dtype={'building_id': int, 'parcel_id': int})
    df['res_sqft_per_unit'] = df['residential_sqft'] / df['residential_units']
    df['res_sqft_per_unit'][df['res_sqft_per_unit'] == np.inf] = 0
    return df

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

# Table of Rental Data from  Craigslist, bayarea_urbansim added by Arezoo

# @orca.table(cache=True)
# def craigslist():
#     df = pd.read_csv(
#         d + 'craigslist.csv',
#         index_col='pid', dtype={'pid': int})
#     return df

@orca.table(cache=True)
def rentals():
    df = pd.read_csv(
        d + 'craigslist.csv',
        index_col='pid', dtype={'pid': int, 'rent': float})
    return df

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

# Tables synthesized by Max Gardner

@orca.table(cache=True)
def units():
    df = pd.read_csv(
        d + 'units.csv',
        index_col='unit_id', dtype={'unit_id': int, 'building_id': int})
    return df

@orca.table(cache=True)
def households():
    try:
        df = pd.read_csv(
            d + 'households.csv',
            index_col='household_id', dtype={
                'household_id': int, 'block_group_id': str, 'state': str,
                'county': str, 'tract': str, 'block_group': str,
                'building_id': int, 'unit_id': int, 'persons': float})
    except:
        df = pd.read_csv(
            d + 'households.csv',
            index_col=0, dtype={
                'block_group_id': str, 'state': str,
                'county': str, 'tract': str, 'block_group': str,
                'building_id': int, 'unit_id': int, 'persons': float})
        df.index.name = 'household_id'
    return df

@orca.table(cache=True)
def persons():
    df = pd.read_csv(
        d + 'persons.csv',
        index_col='person_id', dtype={'person_id': int, 'household_id': int})
    return df

@orca.table(cache=True)
def jobs():

    try:
        df = pd.read_csv(
            d + 'jobs.csv',
            index_col='job_id', dtype={'job_id': int, 'building_id': int})
    except:
        df = pd.read_csv(
            d + 'jobs.csv',
            index_col=0, dtype={'job_id': int, 'building_id': int})
        df.index.name = 'job_id'
    return df

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

# Tables from Sam Blanchard
@orca.table(cache=True)
def establishments():
    df = pd.read_csv(
        d + 'establishments.csv',
        index_col='establishment_id', dtype={
            'establishment_id': int, 'building_id': int, 'primary_id': int})
    return df

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

# zones table
@orca.table(cache=True)
def zones():
    df = pd.read_csv(
        d + 'zones.csv', index_col='zone_id',
        dtype={'zone_id': int})
    return df

############################################################
# skims

# Tables from Emma
@orca.table('skims', cache=True)
def skims():
    df = pd.read_csv(d + 'skims.csv', index_col=0)
    return df

@orca.table(cache=True)
def beam_drive_skims():
    """
    Load BEAM skims, convert travel time to minutes
    """
    df = pd.read_csv(
        d + 'smart-1Apr2019-sc-b-lt-2025-20.skimsExcerpt.csv.gz')

    # morning peak
    df = df[df['period'] == 'AM']
    assert len(df) == 2114116
    df = df.rename(
        columns={'origTaz': 'from_zone_id', 'destTaz': 'to_zone_id'})
    df = df.set_index(['from_zone_id', 'to_zone_id'])

    # seconds to minutes
    df['gen_tt_CAR'] = df['generalizedTimeInS'] / 60
    return df

@orca.table(cache=True)
def beam_skims():
    """
    Load BEAM skims, convert travel time to minutes
    """
    df = pd.read_csv(
        d + 'smart-1Apr2019-sc-b-lt-2025-20.skims.csv.gz')

    df.rename(columns={
        'generalizedCost': 'gen_cost', 'origTaz': 'from_zone_id',
        'destTaz': 'to_zone_id'}, inplace=True)

    # seconds to minutes
    df['gen_tt'] = df['generalizedTimeInS'] / 60

    return df

############################################################
# Broadcasts, a.k.a. merge relationships

orca.broadcast(
    'parcels', 'buildings', cast_index=True, onto_on='parcel_id')
orca.broadcast(
    'buildings', 'units', cast_index=True, onto_on='building_id')
orca.broadcast(
    'units', 'households', cast_index=True, onto_on='unit_id')
orca.broadcast(
    'households', 'persons', cast_index=True, onto_on='household_id')
orca.broadcast(
    'buildings', 'jobs', cast_index=True, onto_on='building_id')
orca.broadcast(
    'buildings', 'establishments', cast_index=True, onto_on='building_id')
orca.broadcast(
    'nodeswalk', 'parcels', cast_index=True, onto_on='node_id_walk')
orca.broadcast(
    'nodeswalk', 'rentals', cast_index=True, onto_on='node_id_walk')
orca.broadcast(
    'nodessmall', 'rentals', cast_index=True, onto_on='node_id_small')
orca.broadcast(
    'nodessmall', 'parcels', cast_index=True, onto_on='node_id_small')
# orca.broadcast(
#     'nodesbeam', 'parcels', cast_index=True, onto_on='node_id_beam')


import numpy as np
import cv2
import os

class Transformer(object):
    def __init__(self, input_dir=None, output_dir=None, data=None, image_size=None):
        self.input_dir = input_dir
        self.output_dir = output_dir
        self.image_name = data.keys()
        self.data = data
        self.image_size = tuple(image_size)
        self.dist_vertices = np.asarray(
            [[0, 0], [image_size[0], 0], [image_size[0], image_size[1]], [0, image_size[1]]])
        # transformed images
        self.image_set = dict()

    def transform(self):
        for image_name in self.image_name:
            image_path = os.path.join(self.input_dir, image_name)
            image = cv2.imread(image_path)
            src_vertices = np.asarray(self.data[image_name])
            M = cv2.getPerspectiveTransform(
                src_vertices.astype(np.float32), self.dist_vertices.astype(np.float32))
            image = cv2.warpPerspective(
                image, M, self.image_size).astype(np.uint8)
            self.image_set.setdefault(image_name, image)

    def save_images(self):
        for image_name in self.image_name:
            out_name = image_name.split(".")[0] + "_out.jpg"
            out_path = os.path.join(self.output_dir, out_name)
            out_image = self.image_set[image_name]
            cv2.imwrite(out_path, out_image)


#!/usr/bin/env python3
from fn_control_prediction import fn_dynamics
from cem_vanilla import fn_control_prediction
from convert_keras_frozen import load_default_frozen_model
import tool_statedata
import rospy
import sys
import numpy as np
"""
roslaunch ur_robot_driver ur5_bringup.launch robot_ip:=192.168.11.3 limited:=false
roslaunch ur5_moveit_config ur5_moveit_planning_execution.launch sim:=false limited:=false
rosrun rqt_controller_manager rqt_controller_manager
rosrun plotjuggler plotjuggler
"""

if __name__ == '__main__':
    rospy.init_node("exp_simulation", disable_signals=True)
    tool_sim  = tool_statedata.datacontroller("simulation")

    # modelname = "Model_8-2_E25"
    # modelname = "Model_8-2_E25_noright"
    # modelname = "Model_8_V6_Combinations_VV1"
    # modelname = "Model_9_beamwall"
    # modelname = "Model10_easy"
    # modelname = "Model10-4_large_noresults"
    # modelname = "Model10-2_result_2"
    # modelname = "Model10-3_result_2_3"
    # modelname = "LSTM_V9_Combinations_2"

    # modelname = "Model_11_V9_Combinations_2"
    modelname = "FCN_Model_11_V9_Combinations_2"
    model = load_default_frozen_model(modelname)

    freq = 20
    p_tgt = 0.0
    x_target = np.array([0,0])
    x_pred = [-1.0, 0, -0.5]
    u = 0

    for i in range(freq*10):
        t_start = rospy.get_time()
        p = x_pred[0]
        dp = x_pred[1]
        th = x_pred[2]
        dth = u

        x_current = np.array([p,dp,th])
        u_next_step, reward = fn_control_prediction(model, x_current, x_target)

        P = 0.1
        D = 0.3
        u_PD = P*(p-p_tgt) + D*dp

        # u = u_PD
        u = u_next_step

        x_pred = fn_dynamics(model, np.expand_dims(x_history,axis=0))[0]

        tool_sim.dataAppend(p, dp, th, dth, u)
        print(rospy.get_time() - t_start)

    tool_sim.dataPlot()
    # tool_sim.detaSave("result_sim")
    sys.exit()


num = int(input("Enter a number : "))

sqr_list = []

for i in range(1, num + 1):
    sqr_list.append(i * i)

print("The list of squares upto", num, "are ", end="")

for sqrs in sqr_list:
    print(sqrs, end=" ")
print()

Part 2 /17 List of Squares/list_of_squares.py


import pygame
import random
import sys
import time

Max_X = 1920
Max_Y = 1080
Max_Snow = 100
Snow_Size = 64 # compress image

class Snow():
    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.speed = random.randint(1, 3) #random speed 1-3
        self.img_num = random.randint(1, 4) #& files
        self.image_filename = "snow" + str(self.img_num) + ".png" #name files
        self.image = pygame.image.load(self.image_filename).convert_alpha() #not bg, clean image convert.alpha
        self.image = pygame.transform.scale(self.image, (Snow_Size, Snow_Size))

    def move_snow(self): # speed snow
        self.y = self.y + self.speed
        if self.y > Max_Y:
            self.y = (0 - Snow_Size)

        i = random.randint(1, 3)
        if i == 1: #move right
            self.x += 1
            if self.x > Max_X:
                self.x = (0 - Snow_Size)
        elif i == 2: #move left
            self.x -= 1
            if self.x < (0 - Snow_Size):
                self.x = Max_X
    def draw_snow(self):
        screen.blit(self.image, (self.x, self.y))

def initialize_snow(Max_Snow, snowfall):
    for i in range (0, Max_Snow):
        xx = random.randint(0, Max_X)
        yy = random.randint(0, Max_Y)
        snowfall.append(Snow(xx, yy)) #coordinate snow

def check_for_exit(): # exit push button
    for event in pygame.event.get():
        if event.type == pygame.KEYDOWN:
            sys.exit()

#--------------------Main------------------

pygame.init()
screen = pygame.display.set_mode((Max_X, Max_Y), pygame.FULLSCREEN) # mode fullscreen

bg_color = (0, 0, 0)
snowfall = [] # storage
initialize_snow(Max_Snow, snowfall) # count snow

while True:
    screen.fill(bg_color) # correct bg
    check_for_exit()
    for i in snowfall:
        i.move_snow()
        i.draw_snow()
    time.sleep(0.0020) #delay
    pygame.display.flip() # draw snowfall


# -*- coding: utf-8 -*-
"""
Created on Fri Jun 30 19:56:13 2023

@author: AMIT
"""
# Smoothening of Images - Apply Blur
import cv2
import numpy as np

img = cv2.imread('E:/OpenCV/child.jpg')
# kernel matrix used for smoothening of images - (3,3),(5,5),(7,7),(9,9)
kernel = np.ones((7,7), np.float32) / 49
print(kernel)

# Filter 2D
smoothed = cv2.filter2D(img, -1, kernel)

# Gaussian Blur
# kernel shape - (3,3),(5,5),(7,7),(9,9)
# sigmaX = std. deviation
g_blur = cv2.GaussianBlur(img, (5,5), 1)

# Median Blur
m_blur = cv2.medianBlur(img, 9)

# Average Blur
avg_blur = cv2.blur(img, (7,7))

cv2.imshow('Image', img)
cv2.imshow('blur', smoothed)
cv2.imshow('g_blur', g_blur)
cv2.imshow('avg_blur', avg_blur)
cv2.imshow('m_blur', m_blur)

# press any to close of frame
cv2.waitKey(0)
# Close all windows
cv2.destroyAllWindows()


#130. Surrounded Regions, Time -O(n) where n -> total values
class Solution:
    def solve(self, board: List[List[str]]) -> None:
        """
        Do not return anything, modify board in-place instead.
        """
        def dfs(i,j,board):
            if i < 0 or i >= len(board) or j < 0 or j >= len(board[i]):
                return
            if 0<=i<len(board) and 0<=j<len(board[i]) and board[i][j] == "O":
                board[i][j] = 'E'
                dfs(i,j-1,board)
                dfs(i,j+1,board)
                dfs(i-1,j,board)
                dfs(i+1,j,board)
        for i in range(len(board)):
            for j in range(len(board[i])):
                if (i in [0,len(board)-1] or j in [0,len(board[i])-1]) and board[i][j] == 'O':
                    dfs(i,j,board)
        for i in range(len(board)):
            for j in range(len(board[i])):
                if board[i][j] == 'E':
                    board[i][j] = 'O'
                else:
                    board[i][j] = 'X'


import matplotlib as mpl

mpl_ver = mpl.__version__.split('.')
if int(mpl_ver[0]) < 2 and int(mpl_ver[1]) < 9:
    colors = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd', '#8c564b', '#e377c2', '#7f7f7f', '#bcbd22', '#17becf']
else:
    colors = ['C{:d}'.format(d) for d in range(10)]

colors = ['#e41a1c',
          '#377eb8',
          '#4daf4a',
          '#984ea3',
          '#ff7f00',
          '#ffff33',
          '#a65628',
          '#f781bf']

colors = ['#1b9e77', '#d95f02', '#e7298a', '#e6ab02','#a6761d','#666666', '#7570b3', '#66a61e']

colors = ['maroon', 'goldenrod', 'mediumseagreen', 'royalblue'] + colors

class Constants():
    def __init__(self):
        self.density = None
        self.fontsize = None
        self.medsize = None
        self.bigfontsize = None
        self.fontname = None
        self.full_font_black = None
        self.folder = None
        self.std_leg_width = None
        self.filetype = None
        self.lettering = None
        self.cbar_height = None

class AnimationConstants(Constants):
    def __init__(self):
        self.density = 220
        self.fontsize = 16
        self.medfontsize = 20
        self.bigfontsize = 32
        self.fontname = 'Computer Modern Sans Serif'
        self.fontfamily = 'sans-serif'
        self.full_font_black = '{:d}p,{:s},black'.format(self.fontsize, self.fontname)
        self.folder = 'forward'
        self.std_leg_width = 1.5
        self.filetype = 'mp4'
        self.lettering = ['a', 'b', 'c', 'd', 'e', 'f']
        self.cbar_height = 0.2

class PrintedConstants(Constants):
    def __init__(self):
        self.density = 300
        self.fontsize = 8
        self.medfontsize = 12
        self.bigfontsize = 16
        self.fontname = 'Times-Roman'
        self.fontfamily = 'serif'
        self.full_font_black = '{:d}p,{:s},black'.format(self.fontsize, self.fontname)
        self.folder = 'pub'
        self.std_leg_width = 1.5
        self.filetype = 'png'
        self.lettering = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i']
        self.cbar_height = 0.1

class ScreenConstants(Constants):
    def __init__(self):
        self.density = 450
        self.fontsize = 12
        self.medfontsize = 14
        self.bigfontsize = 20
        self.fontname = 'Computer Modern Sans Serif'
        self.fontfamily = 'sans-serif'
        self.full_font_black = '{:d}p,{:s},black'.format(self.fontsize, self.fontname)
        self.folder = 'slideshow'
        self.std_leg_width = 1.5
        self.filetype = 'pdf'
        self.lettering = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i']
        self.cbar_height = 0.1

def load_constants(cname, **kwargs):
    print('loading constants')
    if cname in ['serif', 'print', 'pub']:
        constants = PrintedConstants()
    elif cname in ['animation', 'movie']:
        constants = AnimationConstants()
    elif cname in ['sans', 'sans-serif', 'screen']:
        constants = ScreenConstants()
    else:
        raise ValueError('bad cname')

    for key, value in kwargs.items():
        setattr(constants, key, value)

    if False:
        try:
            mpl.use('macosx')
        except:
            if constants.filetype == 'pdf':
                mpl.use('PDF')
            elif constants.filetype == 'png':
                mpl.use('agg')
            else:
                mpl.use('TKagg')

    mpl.rc('text', usetex=True)
    mpl.rc('font', **{'family': constants.fontfamily, constants.fontfamily: constants.fontname})
    if constants.fontfamily == 'sans-serif':
        mpl.rc('text.latex', preamble=r'\usepackage{color} \usepackage{sfmath} \renewcommand{\familydefault}{\sfdefault} \usepackage{cmbright}')
    else:
        mpl.rc('text.latex', preamble=r'\usepackage{color}')
    try:
        mpl.rcParams['patch.force_edgecolor'] = True
    except KeyError:
        # This is going to happen with some old versions of mpl
        pass

    return constants


import ROOT
from ROOT import TMVA

outputFile = ROOT.TFile("TMVA.root","RECREATE")
factory = TMVA.Factory("tmvaTest", outputFile, "")
dataloader = TMVA.DataLoader("dataset")

# trainingFile = ROOT.TFile("../generate/trainingData.root")
# testFile = ROOT.TFile("../generate/testData.root")
inputFile = ROOT.TFile("../data/TrainingFile.root")

# get the TTree objects from the input files

sig = inputFile.Get("TreeQ")
bkg = inputFile.Get("TreeG")
nSig = sig.GetEntries()
nBkg = bkg.GetEntries()

# global event weights (see below for setting event-wise weights)

sigWeight = 1.0
bkgWeight = 1.0

# dataloader.AddSignalTree(sigTest, sigWeight, TMVA.Types.kTesting)
# dataloader.AddBackgroundTree(bkgTest, bkgWeight, TMVA.Types.kTesting)

dataloader.AddVariable("qpTD", 'F')
dataloader.AddVariable("qSigma2", 'F')
dataloader.AddVariable("qMulti", 'F')

dataloader.AddSignalTree(sig, sigWeight)
dataloader.AddBackgroundTree(bkg, bkgWeight)
cut1 = ROOT.TCut("JetPt > 30 && JetPt < 100")
cut2 = ROOT.TCut("")
dataloader.PrepareTrainingAndTestTree(cut1,"SplitMode=Random:!V")

# Book MVA methods (see TMVA manual).
#factory.BookMethod(dataloader,TMVA.Types.kMLP, "MLP", "H:!V:HiddenLayers=3");
factory.BookMethod(dataloader,TMVA.Types.kBDT, "BDT200", "NTrees=200:BoostType=AdaBoost")
# Train, test and evaluate all methods

factory.TrainAllMethods()
factory.TestAllMethods()
factory.EvaluateAllMethods()

outputFile.Close()
print(' TMVAanalyses is done! ')


import sys
sys.path.append("/mnt/c/Users/Harris/Google Drive/LHCb/bddk/Analysis_2021/")
from essentials import *
import rootutils
# from Analysis_2021.essentials import get_shapes_bkg, get_free_shapes, save_png, MakeSWeights
import glob as glob
# from Analysis_2021 import rootutils
from uncertainties import ufloat
from uncertainties import covariance_matrix, correlation_matrix
from uncertainties import unumpy, ufloat_fromstr
import math as math
import pandas
import numpy as np
RDF = ROOT.ROOT.RDataFrame

B0_ID = 511
Bp_ID = 521
Bs_ID = 531

Dp_ID = 411
D0_ID = 421
Ds_ID = 431

Dpst_ID = 413
D0st_ID = 423
Dsst_ID = 433

k_ID = 321
pi_ID = 211
kst0_ID = 313

def build_truth_base(b_spec, c1_spec, c2_spec, norm_flag = 0, st_flag = 0):
    truth_MC = f"abs(B_TRUEID) == {b_spec}"
    truth_MC = f"{truth_MC} && abs(D1H1_TRUEID) == {k_ID}"
    truth_MC = f"{truth_MC} && abs(D1H2_TRUEID) == {pi_ID}"
    if c1_spec == Dp_ID:
        truth_MC = f"{truth_MC} && abs(D1H3_TRUEID) == {pi_ID}"
    truth_MC = f"{truth_MC} && abs(D2H1_TRUEID) == {k_ID}"
    truth_MC = f"{truth_MC} && abs(D2H2_TRUEID) == {pi_ID}"
    if c2_spec == Dp_ID:
        truth_MC = f"{truth_MC} && abs(D2H3_TRUEID) == {pi_ID}"
    if norm_flag == 1:
        truth_MC = f"{truth_MC} && abs(D2H4_TRUEID) == {pi_ID}"
        truth_MC = f"{truth_MC} && abs(K_TRUEID) == {k_ID}"
    if norm_flag == 0:
        truth_MC = f"{truth_MC} && abs(KSTH1_TRUEID) == {k_ID}"
        truth_MC = f"{truth_MC} && abs(KSTH2_TRUEID) == {pi_ID}"
    return truth_MC

def build_truth_mom(b_spec, c1_spec, c2_spec, norm_flag = 0, st_flag = 0):
    truth_MC = f"abs(D1_TRUEID) == {c1_spec}"
    truth_MC = f"{truth_MC} && abs(D2_TRUEID) == {c2_spec}"
    if norm_flag == 0:
        truth_MC = f"{truth_MC} && abs(KST_TRUEID) == {kst0_ID}"
    return truth_MC

def get_tmc_ratio(path, mc_spec, data_ws, plot_flag = False, dict_flag = False, snap_sideband_flag = False, fit_sideband_flag = True):

    sig_files = glob.glob(f"/mnt/c/Users/Harris/Desktop/rootfiles/2022_filtered_mc/*/sig_d/{path}.root")
    sb_files = glob.glob(f"/mnt/c/Users/Harris/Desktop/rootfiles/2022_filtered_mc/*/sb_d/{path}.root")

    tc_sig = ROOT.TChain(f"DecayTreeTuple")
    tc_sb = ROOT.TChain(f"DecayTreeTuple")

    for file in sig_files:
        tc_sig.Add(file)

    for file in sb_files:
        tc_sb.Add(file)

    rdf_sig = RDF(tc_sig)
    rdf_sb = RDF(tc_sb)

    truth_base_string = build_truth_base(B0_ID, Dp_ID, Dp_ID)
    # truth_mom_string = build_truth_mom(B0_ID, Dp_ID, Dp_ID)

    # print(truth_base_string)
    # print(truth_mom_string)
    #({truth_mom_string}) &&
    # ({truth_mom_string}) &&

    rdf_tmc_sig = rdf_sig.Filter(f"({truth_base_string})", "truth_all")
    rdf_tmc_sb = rdf_sb.Filter(f"({truth_base_string})", "truth_all")

    tmc_sig_count = rdf_tmc_sig.Count().GetValue()
    tmc_sb_count = rdf_tmc_sb.Count().GetValue()
    mean_start = rdf_tmc_sb.Mean("B_M").GetValue()

    if snap_sideband_flag:
        clist = rdf_tmc_sb.GetColumnNames()
        rdf_tmc_sb_snap = rdf_tmc_sb.Snapshot(f"DecayTreeTuple", f"sb_tmc1_files/{mc_spec}.root", clist)
        print(f"finished snapshot for tmc1 mc")

    if fit_sideband_flag:

        sbws = ROOT.RooWorkspace("sbws")
        bmin = 5000
        bmax = 5400
        sbws.factory(f"B_M[{bmin},{bmax}]")
        b_m = sbws.var("B_M")
        data_args = ROOT.RooArgSet(b_m)
        sbfile = ROOT.TFile(f"sb_tmc1_files/{mc_spec}.root")
        sbtree = sbfile.Get("DecayTreeTuple")
        lmbdata = ROOT.RooDataSet(f"{mc_spec}_lmbsbmc", f"{mc_spec}_lmbsbmc", sbtree, data_args)

        sbws.Import(lmbdata)
        sbws.factory(f"RooKeysPdf::{mc_spec}_lmb_fit(B_M, {mc_spec}_lmbsbmc)")
        lmbfit = sbws.pdf(f"{mc_spec}_lmb_fit")
        data_ws.Import(lmbfit)

    if plot_flag:

        hist_bm_sig = rdf_tmc_sig.Histo1D((f"bm_sig", f"bm_sig", 100, 4800, 5600), 'B_M')
        hist_bm_sb = rdf_tmc_sb.Histo1D((f"bm_sb", f"bm_sb", 100, 4800, 5600), 'B_M')

        c1 = ROOT.TCanvas("c1","c1")
        hist_bm_sig.Draw()
        save_png(c1, f"test/{path}", f"{path}_sig", rpflag = 0)

        c2 = ROOT.TCanvas("c2","c2")
        hist_bm_sb.Draw()
        save_png(c2, f"test/{path}", f"{path}_sb", rpflag = 0)

        frame = b_m.frame(ROOT.RooFit.Title(f"{mc_spec}_mcsb_spectrum"), ROOT.RooFit.Bins(40))

        lmbdata.plotOn(frame, ROOT.RooFit.Name("data"))
        lmbfit.plotOn(frame, ROOT.RooFit.LineStyle(ROOT.kSolid), ROOT.RooFit.LineColor(ROOT.kRed), ROOT.RooFit.Name("pdf"))

        p = rootutils.residualPlot()
        p.pt.cd()
        xaxis = frame.GetXaxis()
        xaxis.SetTickLength(0)
        xaxis.SetNdivisions(0)
        xaxis.SetLabelSize(0)

        frame.Draw()
        p.pb.cd()

        hpull = frame.pullHist(f"data", f"pdf")
        pull = b_m.frame()
        pull.addPlotable(hpull, "P")
        pull.SetLineWidth(ROOT.gStyle.GetFrameLineWidth())
        pull.GetXaxis().SetLabelSize(ROOT.gStyle.GetLabelSize() * p.blabelratio)
        pull.GetYaxis().SetLabelSize(
            ROOT.gStyle.GetLabelSize("Y") * (1 + p.padratio) / (2 * p.padratio)
        )
        pull.GetXaxis().SetTitleSize(ROOT.gStyle.GetTitleSize() * p.blabelratio)
        pull.GetYaxis().SetTitleSize(ROOT.gStyle.GetTitleSize() * p.blabelratio)
        pull.GetYaxis().SetTitleOffset(ROOT.gStyle.GetTitleOffset() / p.blabelratio)
        pull.GetXaxis().SetTickLength(ROOT.gStyle.GetTickLength() * p.blabelratio)
        pull.GetYaxis().SetTitle("Pull")
        pull.GetYaxis().SetNdivisions(202, False)
        pull.GetYaxis().SetRangeUser(-3, 3)
        pull.GetYaxis().CenterTitle()
        pull.Draw("AP")

        # pull.GetXaxis().SetTitle(f" m({title} {varoi}) [MeV]")
        # frame.GetXaxis().SetTitle(f" m({title}) [MeV]")

        save_png(p, f"sb_keys_fit_tests", f"{mc_spec}_sb", rpflag = 1)
    if dict_flag:
        dict = {
            "Scheme ID" : path,
            "TMC in Sig": tmc_sig_count,
            "TMC in sb": tmc_sb_count,
        }

        folder_path = '/mnt/c/Users/Harris/Google Drive/LHCb/bddk/Analysis_2021/charmless_bkg_estimation/charmless_ratio'
        if not os.path.exists(folder_path):
            os.makedirs(folder_path)

        df = pandas.DataFrame(dict , index=[0])
        df.to_csv(f"{folder_path}/{path}_charm_txt")

    return (tmc_sb_count/tmc_sig_count, mean_start)

def get_shape(dws, data_spec, mc_spec, run_name, varoi):

    nn_ws_base = ROOT.TFile(f"/mnt/c/Users/Harris/Google Drive/LHCb/bddk/Analysis_2021/data_fit/data_files/{data_spec}_{run_name}_{varoi}.root")
    nn_ws = nn_ws_base.Get(run_name)

    if mc_spec == "Z_m_p_01":
        file = "01_Z_m_p_11198006"
    if mc_spec == "Z_m_p_0203":
        file = "02_Z_m_p_11198400"
    # Get Signal Yield
    y = nn_ws.var(f"{mc_spec}_yield")
    # Get TMC sb ratio
    tmc_ratio, mean_lmb = get_tmc_ratio(file, mc_spec, dws)
    # Get mean for charmless
    mean_charmless = nn_ws.var(f"mean_{mc_spec}")
    dws.Import(mean_charmless)
    lmb_yield = tmc_ratio*y.getValV()
    dws.factory(f"{mc_spec}_lmb_yield[{lmb_yield}]")
    dws.factory(f"Gaussian::{mc_spec}_charmless_fit(B_M, mean_{mc_spec}, width_{mc_spec}_charmless[5.0,0.1,15.0])")

def build_charmless_fit(run_name, data_spec):

    dws = ROOT.RooWorkspace(run_name)
    bmin = 5000
    bmax = 5400
    dws.factory(f"B_M[{bmin},{bmax}]")

    get_shape(dws, data_spec, "Z_m_p_01", run_name, "B_M")
    get_shape(dws, data_spec, "Z_m_p_0203", run_name, "B_M")
    get_shapes_bkg(data_spec, "Exponential", dws, varoi = "B_M")
    # get_shapes_bkg(data_spec, "Bernstein", dws, varoi = "B_M")

    dws.factory("SUM::Z_m_p_spectrum_all_fit(Z_m_p_01_charmless_yield[500,0,10000]*Z_m_p_01_charmless_fit, \
                                             Z_m_p_01_lmb_yield*Z_m_p_01_lmb_fit, \
                                             Z_m_p_0203_charmless_yield[500,0,10000]*Z_m_p_0203_charmless_fit, \
                                             Z_m_p_0203_lmb_yield*Z_m_p_0203_lmb_fit, \
                                             Z_m_p_bkg_yield[100,0,100000]*Z_m_p_spectrum_bkg)")

    dws.Print()
    b_m = dws.var("B_M")
    data_args = ROOT.RooArgSet(b_m)

    file_list = glob.glob(f"/mnt/c/Users/Harris/Desktop/rootfiles/2022_filtered_data/*/final_sample/{data_spec}.root")
    tchain = ROOT.TChain("DecayTreeTuple_SB")
    for file_name in file_list:
        tchain.Add(file_name)
    #
    data = ROOT.RooDataSet(f"{data_spec}_final_data", f"{data_spec}_final_data", tchain, data_args)
    model = dws.pdf(f"{data_spec}_spectrum_all_fit")
    fit = model.fitTo(data, ROOT.RooFit.PrintLevel(0), ROOT.RooFit.Save())
    dws.Import(fit)
    fitresult = dws.obj(f"fitresult_{data_spec}_spectrum_all_fit_{data_spec}_final_data")

    frame = b_m.frame(ROOT.RooFit.Title(f"{data_spec}_spectrum"), ROOT.RooFit.Bins(40))

    ylist = ["Z_m_p_01_charmless_fit","Z_m_p_01_lmb_fit","Z_m_p_0203_charmless_fit","Z_m_p_0203_lmb_fit"]

    data.plotOn(frame, ROOT.RooFit.Name("data"))
    model.plotOn(frame, ROOT.RooFit.LineStyle(ROOT.kSolid), ROOT.RooFit.LineColor(ROOT.kRed), ROOT.RooFit.Name("pdf"))
    model.plotOn(frame, ROOT.RooFit.Components(f"{data_spec}_spectrum_bkg"), ROOT.RooFit.LineStyle(ROOT.kSolid), ROOT.RooFit.LineColor(ROOT.kGreen), ROOT.RooFit.Name("bkg"))

    color_list = [ROOT.kBlue, ROOT.kViolet, ROOT.kAzure, ROOT.kMagenta, ROOT.kCyan]

    for y, color in zip(ylist, color_list):
        model.plotOn(frame, ROOT.RooFit.Components(y), ROOT.RooFit.LineStyle(ROOT.kDashed), ROOT.RooFit.LineColor(color), ROOT.RooFit.Name(f"fit_{y}"))

    legend = ROOT.TLegend(0.20, 0.575, 0.40, 0.905)
    name_list =  ["01_charmless_fit","01_lmb_fit","0203_charmless_fit","0203_lmb_fit"]

    legend.SetFillStyle(1001)

    d_nyield_bkg = dws.obj(f"{data_spec}_bkg_yield")
    d_nyield_bkg_err = d_nyield_bkg.getPropagatedError(fitresult)

    legend.AddEntry(frame.findObject("data"),"Run 2 Data","ep")
    legend.AddEntry(frame.findObject("pdf"),"Total Fit PDF","l")
    legend.AddEntry(frame.findObject("bkg"),f"Background PDF : {round(d_nyield_bkg.getValV(),1)} #pm {round(d_nyield_bkg_err,1)}","l")
    for y, name in zip(ylist, name_list):
        yname = y.split("_fit")[0]
        print(yname)
        d_nyield = dws.obj(f"{yname}_yield")
        d_nyield_err = d_nyield.getPropagatedError(fitresult)
        legend.AddEntry(frame.findObject(f"fit_{y}"), f"{name}: {round(d_nyield.getValV(),1)} #pm {round(d_nyield_err,1)}", "l")
    legend.SetTextSize(0.020)

    p = rootutils.residualPlot()
    p.pt.cd()
    xaxis = frame.GetXaxis()
    xaxis.SetTickLength(0)
    xaxis.SetNdivisions(0)
    xaxis.SetLabelSize(0)

    frame.Draw()
    legend.Draw()
    p.pb.cd()

    hpull = frame.pullHist(f"data", f"pdf")
    pull = b_m.frame()
    pull.addPlotable(hpull, "P")
    pull.SetLineWidth(ROOT.gStyle.GetFrameLineWidth())
    pull.GetXaxis().SetLabelSize(ROOT.gStyle.GetLabelSize() * p.blabelratio)
    pull.GetYaxis().SetLabelSize(
        ROOT.gStyle.GetLabelSize("Y") * (1 + p.padratio) / (2 * p.padratio)
    )
    pull.GetXaxis().SetTitleSize(ROOT.gStyle.GetTitleSize() * p.blabelratio)
    pull.GetYaxis().SetTitleSize(ROOT.gStyle.GetTitleSize() * p.blabelratio)
    pull.GetYaxis().SetTitleOffset(ROOT.gStyle.GetTitleOffset() / p.blabelratio)
    pull.GetXaxis().SetTickLength(ROOT.gStyle.GetTickLength() * p.blabelratio)
    pull.GetYaxis().SetTitle("Pull")
    pull.GetYaxis().SetNdivisions(202, False)
    pull.GetYaxis().SetRangeUser(-3, 3)
    pull.GetYaxis().CenterTitle()
    pull.Draw("AP")

    # pull.GetXaxis().SetTitle(f" m({title} {varoi}) [MeV]")
    # frame.GetXaxis().SetTitle(f" m({title}) [MeV]")

    save_pdf(p, f"fit_tests", f"{data_spec}_{run_name}_nl", rpflag = 1)

build_charmless_fit("charmlesstest", "Z_m_p")


#!/bin/usr/python
#------------------------------------------------------------------------------
# File Name: download_tools.py
# Author 1: Qingkai Kong
# Author 2: Noah G. Luna
#
# Functions used for various stages of mass file query, downloading and format.
# Do not modify without checking dependies of files within main directory!
#------------------------------------------------------------------------------

import os
from obspy.clients.fdsn import Client
from obspy import UTCDateTime
from obspy.core.event import read_events
from joblib import Parallel, delayed

import pandas as pd
from datetime import datetime
import pickle

from tools import find_regional_station_distance

import glob
import errno

def get_station_pick(filename, stations):
    station_picks = []
    print(filename)
    cat = read_events(filename)
    # loop through the picks and
    for n, event in enumerate(cat):

        if len(event.origins) == 0:
            continue
        elif len(event.origins) > 1:
            print('More origins present %d'%(len(event.origins)))
            origin_info = []

            for ooo in event.origins:
                origin_info += ooo.arrivals
        else:
            origin_info = event.origins[0].arrivals

        for i, pick in enumerate(event.picks):
            waveform_id = pick['waveform_id']
            stnm = waveform_id['station_code']
            channel = waveform_id['channel_code']
            network = waveform_id['network_code']
            onset = pick.onset
            polarity = pick.polarity
            evaluation_mode = pick.evaluation_mode
            evaluation_status = pick.evaluation_status
            phase_time = pick.time
            origin = event.origins[0]
            evtime = origin.time
            evla = origin.latitude
            evlo = origin.longitude
            try:
                stla, stlo, stel = stations[network + '_' + stnm]
            except:
                print('Can not find %s, %s'%(network, stnm))
                continue

            dist = find_regional_station_distance(evla, evlo, stla, stlo)
#            if dist > 30:
#                continue

            mags = event.magnitudes[0]
            mag = mags.mag
            mag_type = mags.magnitude_type
            eventId = event.resource_id.id.split('/')[-1]

            arrival = origin_info[i]
            pick_id = pick.resource_id
            pid = arrival['pick_id']

            if pid != pick_id:

                new_arrival = \
                    [arrival for arrival in origin_info if arrival['pick_id'] == pick_id]
                if len(new_arrival) < 1:
                    continue
                phase = new_arrival[0]['phase']
            else:
                phase = arrival['phase']

            station_picks.append([network, stnm, channel,dist, phase, phase_time, onset, polarity, evaluation_mode, evaluation_status, \
                                    evtime, evla, evlo, mag, mag_type, eventId])

    return station_picks

def get_station_data(row, stations, client):
    t = row.phase_time
    # Original time window for clean & noisy traces
    # These were 5 minute records
    t0 = t - 60
    t1 = t + 240

    net = row.network
    stnm = row.stnm
    comp = row.channel
    onset = row.onset
    status = row.evaluation_status
    mode = row.evaluation_mode
    timestamp = str(int(t.timestamp*1000))
    eventId = row.eventId

    stla, stlo, stel = stations[net + '_' + stnm]
    if onset is None:
        return

    if (comp.lower()[-1] != 'z') & (comp[-1] != '3'):
        #print(comp)
        return

    # Output file directory
    output_file_dir = './mseed'

    # Check if folder exists:
    if not os.path.isdir(output_file_dir):
        os.makedirs(output_file_dir)

    output_name =  net + '_' + stnm + '_' \
          + '_' + onset + '_' + status + '_' + eventId

    output = os.path.join(output_file_dir, output_name)

    # skip for duplicates
    if os.path.exists(output ):
        print("Skiping " + output, comp)
        return

    #if row.onset != 'impulsive':
    #    continue
    try:
        st = client.get_waveforms(net, stnm, '*', channel='BH?,HH?', starttime=t0, endtime=t1)
    except:
        print('No data for %s, %s, %s, %s, %s'%(net, stnm, comp, t0, t1))
        return

    # put the data into the
    for tr in st:
        tr.stats.sac = {}
        tr.stats.sac.evla = row.evla
        tr.stats.sac.evlo = row.evlo
        tr.stats.sac.mag = row.mag
        tr.stats.sac.stla = stla
        tr.stats.sac.stlo = stlo
        comp = tr.stats.channel

        #tr.write(output + '_' + comp + '.SAC', format = 'SAC')
    st.write(output + '.mseed', format = 'MSEED')
    print('writing ' + output)
    dist = row.dist

def get_noise(row, stations, client):
    t = row.phase_time
    # Here we get data 20 minutes and 10 seconds before event
    t0 = row.noise_startt
    # The record will be 8 minutes (480 seconds) long
    t1 = row.noise_endt

    net = row.network
    stnm = row.stnm
    comp = row.channel
    timestamp = str(int(t.timestamp*1000))
    eventId = row.eventId

    stla, stlo, stel = stations[net + '_' + stnm]

    if (comp.lower()[-1] != 'z') & (comp[-1] != '3'):
        #print(comp)
        return

    # Output file directory
    output_file_dir = './mseed_noise'

    # Check if folder exists:
    if not os.path.isdir(output_file_dir):
        try:
            os.makedirs(output_file_dir)
        except OSError as e:
            if e.errno != errno.EEXIST:
                raise

    # File name
    output_name = './mseed_noise/' + 'noise' + '_' + net + '_' + stnm + '_' + eventId
    os.path.join(output_file_dir, output_name)

    # skip for duplicates
    if os.path.exists(output_name):
        print("Skiping " + output_name, comp)
        return

    try:
        st = client.get_waveforms(net, stnm, '*', channel='BH?,HH?', starttime=t0, endtime=t1)
    except:
        print('No data for %s, %s, %s, %s, %s'%(net, stnm, comp, t0, t1))
        return

    # put the data into the
    for tr in st:
        tr.stats.sac = {}
        tr.stats.sac.evla = row.evla
        tr.stats.sac.evlo = row.evlo
        tr.stats.sac.mag = row.mag
        tr.stats.sac.stla = stla
        tr.stats.sac.stlo = stlo
        comp = tr.stats.channel

        #tr.write(output + '_' + comp + '.SAC', format = 'SAC')
    st.write(output_name + '.mseed', format = 'MSEED')
    print('writing ' + output_name)
    dist = row.dist

# Download catalog function
def download_catalog(i, date_range, client, stations, minmag, maxrad):

    t0 = UTCDateTime(date_range[i])
    t1 = UTCDateTime(date_range[i+1])

    for coord in stations:
        station = coord
        stla, stlo, stel = stations[coord]

        output = t0.strftime('%Y-%m-%d') + '_' + t1.strftime('%Y-%m-%d') + '_' + station + '.quakeml'

        # if file exist, we just skip it
        if os.path.exists(output):
            return

        try:
            cat = client.get_events(starttime=t0, endtime=t1, latitude=stla, longitude=stlo,\
                maxradius=maxrad, minmagnitude=minmag,\
                includearrivals = True, filename = output)

        except Exception as e:
            print('There was an error in obtaining data for station', station.upper())
            print('Time stamp: ', t0.strftime('%Y-%m-%d') + '_' + t1.strftime('%Y-%m-%d') +' will be skipped\n')


from sqlite3_helper import sqlite3_create, sqlite3_run
import os

sqlite3_file = 'drivers.sqlite3'
sqlite3_drivers_schema = 'drivers.sql'
sqlite3_trips_schema = 'trips.sql'

def run_query(query):
        sqlite3_output = sqlite3_run(sqlite3_file, query)

        return sqlite3_output

def sqlite3_check():
    schema_list=[sqlite3_drivers_schema, sqlite3_trips_schema]

    for schema in schema_list:
        if os.path.exists(schema) and os.path.isfile(schema):
            with open (schema) as file:
                sqlite3_create(sqlite3_file, file.read())
        else:
            output = f'Error: Required file "{schema}" is not exist'

            return output

def driver(driver_name, driver_family_name, driver_plate_number):
    if sqlite3_check() == None:
        if run_query(f'SELECT id FROM drivers WHERE name = "{driver_name}" AND family_name = "{driver_family_name}" AND plate_number = "{driver_plate_number}";') == []:

            add_record = f""" INSERT INTO drivers (name, family_name, plate_number) VALUES (
                            "{driver_name}", "{driver_family_name}", "{driver_plate_number}"
                            );"""
            run_query(add_record)
    else:
        output = sqlite3_check()
        return output

def trips(driver_name, driver_family_name, driver_plate_number):
    if sqlite3_check() == None:
        if run_query(f'SELECT id FROM drivers WHERE name = "{driver_name}" AND family_name = "{driver_family_name}" AND plate_number = "{driver_plate_number}";') == []:

            add_record = f""" INSERT INTO drivers (name, family_name, plate_number) VALUES (
                            "{driver_name}", "{driver_family_name}", "{driver_plate_number}"
                            );"""
            run_query(add_record)
    else:
        output = sqlite3_check()
        return output

def main():

    driver_name = 'Amedd'
    driver_family_name = 'Hame'
    driver_plate_number = '125 ABC'

    driver(driver_name, driver_family_name, driver_plate_number)

if __name__ == "__main__":
    main()


n,m=map(int,input().split())
arr=list(map(int,input().split()))
summax = 0

for i in range (n-2):
    for j in range (i+1,n-1):
        for k in range (j+1,n):
            if arr[i] + arr[j] + arr[k]<=m:
                if summax >= arr[i] + arr[j] + arr[k]:
                    continue
                else:
                    summax = arr[i] + arr[j] + arr[k]
            else:
                continue

print(summax)


n = int(input())
call = []
for _ in range(n):
    call.append(list(input().split()))

nums_list = []

for i in range(1, 10):
    for j in range(1, 10):
        if j != i:
            for k in range(1, 10):
                if k != i and k != j:
                    nums_list.append(str(i)+str(j)+str(k))

def check_nums(num):
    cnt = 0
    for i in range(n):
        strike = 0
        ball = 0
        for j in range(3):
            if num[j] == call[i][0][j]:
                strike += 1
        for j in range(3):
            if num[j] != call[i][0][j] and num[j] in call[i][0]:
                ball += 1
        if strike == int(call[i][1]) and ball == int(call[i][2]):
            cnt += 1

    if cnt == n:
        return True
    else:
        return False

result = 0
for i in nums_list:
    if check_nums(i):
        result += 1

print(result)


import gc

import numpy as np
import torch.nn as nn
import torch

class EffectExplainer(nn.Module):
    def __init__(
            self,
            model
    ):
        super(EffectExplainer, self).__init__()
        self.model = model

    def run_explain(self, feat, adj):
        with torch.no_grad():
            self.model.using_attention_explanation()
            _ = self.model(feat, adj)
        edge_weight_matrix = self.model.attention_weight
        return edge_weight_matrix


import random

@profile
def SCAN(a):
    n = len(a)
    for i in range(n-1):
        for j in range(i + 1, n):
            if a[i] == a[j]:
                print(i, " - ", j)
                print('SCAN Duplicate Element', a[i])
                return
            else:
                continue
    return

@profile
def STOR(a):
    b = []
    for i in a:
        if i not in b:
            b.append(i)
        else:
            print('STOR Duplicate Element', i)
            print('finished', b)

def generateArray():
    arr = list(range(1, 1001))
    arr.append(random.randint(0, 1000))

    random.shuffle(arr)
    return arr

array_a = generateArray()

SCAN(array_a)
STOR(array_a)


import os
from collections import OrderedDict

import numpy as np
import torch
from tqdm import tqdm

from confidnet.learners.learner import AbstractLeaner
from confidnet.utils import misc
from confidnet.utils.logger import get_logger
from confidnet.utils.metrics import Metrics

LOGGER = get_logger(__name__, level="DEBUG")

class SelfConfidLearner(AbstractLeaner):
    def __init__(self, config_args, train_loader, val_loader, test_loader, start_epoch, device):
        super().__init__(config_args, train_loader, val_loader, test_loader, start_epoch, device)
        self.freeze_layers()
        self.disable_bn(verbose=True)
        if self.config_args["model"].get("uncertainty", None):
            self.disable_dropout(verbose=True)

    def train(self, epoch):
        self.model.train()
        self.disable_bn()
        if self.config_args["model"].get("uncertainty", None):
            self.disable_dropout()
        metrics = Metrics(
            self.metrics, self.prod_train_len, self.num_classes
        )
        loss, confid_loss = 0, 0
        len_steps, len_data = 0, 0

        # Training loop
        loop = tqdm(self.train_loader)
        for batch_id, (data, target) in enumerate(loop):
            data, target = data.to(self.device), target.to(self.device)
            self.optimizer.zero_grad()
            output = self.model(data)

            # Potential temperature scaling
            if self.temperature:
                output = list(output)
                output[0] = output[0] / self.temperature
                output = tuple(output)

            if self.task == "classification":
                current_loss = self.criterion(output, target)
            elif self.task == "segmentation":
                current_loss = self.criterion(output, target.squeeze(dim=1))
            current_loss.backward()
            loss += current_loss
            self.optimizer.step()
            if self.task == "classification":
                len_steps += len(data)
                len_data = len_steps
            elif self.task == "segmentation":
                len_steps += len(data) * np.prod(data.shape[-2:])
                len_data += len(data)

            # Update metrics
            pred = output[0].argmax(dim=1, keepdim=True)
            confidence = torch.sigmoid(output[1])
            metrics.update(pred, target, confidence)

            # Update the average loss
            loop.set_description(f"Epoch {epoch}/{self.nb_epochs}")
            loop.set_postfix(
                OrderedDict(
                    {
                        "loss_confid": f"{(loss / len_data):05.3e}",
                        "acc": f"{(metrics.accuracy / len_steps):05.2%}",
                    }
                )
            )
            loop.update()

        # Eval on epoch end
        scores = metrics.get_scores(split="train")
        logs_dict = OrderedDict(
            {
                "epoch": {"value": epoch, "string": f"{epoch:03}"},
                "lr": {
                    "value": self.optimizer.param_groups[0]["lr"],
                    "string": f"{self.optimizer.param_groups[0]['lr']:05.1e}",
                },
                "train/loss_confid": {
                    "value": loss / len_data,
                    "string": f"{(loss / len_data):05.4e}",
                },
            }
        )
        for s in scores:
            logs_dict[s] = scores[s]

        # Val scores
        val_losses, scores_val = self.evaluate(self.val_loader, self.prod_val_len, split="val")
        logs_dict["val/loss_confid"] = {
            "value": val_losses["loss_confid"].item() / self.nsamples_val,
            "string": f"{(val_losses['loss_confid'].item() / self.nsamples_val):05.4e}",
        }
        for sv in scores_val:
            logs_dict[sv] = scores_val[sv]

        # Test scores
        test_losses, scores_test = self.evaluate(self.test_loader, self.prod_test_len, split="test")
        logs_dict["test/loss_confid"] = {
            "value": test_losses["loss_confid"].item() / self.nsamples_test,
            "string": f"{(test_losses['loss_confid'].item() / self.nsamples_test):05.4e}",
        }
        for st in scores_test:
            logs_dict[st] = scores_test[st]

        # Print metrics
        misc.print_dict(logs_dict)

        # Save the model checkpoint
        self.save_checkpoint(epoch)

        # CSV logging
        misc.csv_writter(path=self.output_folder / "logs.csv", dic=OrderedDict(logs_dict))

        # Tensorboard logging
        self.save_tb(logs_dict)

        # Scheduler step
        if self.scheduler:
            self.scheduler.step()

    def evaluate(self, dloader, len_dataset, split="test", verbose=False, **args):
        self.model.eval()
        metrics = Metrics(self.metrics, len_dataset, self.num_classes)
        loss = 0

        # Evaluation loop
        loop = tqdm(dloader, disable=not verbose)
        for batch_id, (data, target) in enumerate(loop):
            data, target = data.to(self.device), target.to(self.device)

            with torch.no_grad():
                output = self.model(data)
                if self.task == "classification":
                    loss += self.criterion(output, target)
                elif self.task == "segmentation":
                    loss += self.criterion(output, target.squeeze(dim=1))
                # Update metrics
                pred = output[0].argmax(dim=1, keepdim=True)
                confidence = torch.sigmoid(output[1])
                metrics.update(pred, target, confidence)

        scores = metrics.get_scores(split=split)
        losses = {"loss_confid": loss}
        return losses, scores

    def load_checkpoint(self, state_dict, uncertainty_state_dict=None, strict=True):
        if not uncertainty_state_dict:
            self.model.load_state_dict(state_dict, strict=strict)
        else:
            self.model.pred_network.load_state_dict(state_dict, strict=strict)

            # 1. filter out unnecessary keys
            if self.task == "classification":
                state_dict = {
                    k: v
                    for k, v in uncertainty_state_dict.items()
                    if k not in ["fc2.weight", "fc2.bias"]
                }
            if self.task == "segmentation":
                state_dict = {
                    k: v
                    for k, v in uncertainty_state_dict.items()
                    if k
                    not in [
                        "up1.conv2.cbr_unit.0.weight",
                        "up1.conv2.cbr_unit.0.bias",
                        "up1.conv2.cbr_unit.1.weight",
                        "up1.conv2.cbr_unit.1.bias",
                        "up1.conv2.cbr_unit.1.running_mean",
                        "up1.conv2.cbr_unit.1.running_var",
                    ]
                }
            # 2. overwrite entries in the existing state dict
            self.model.uncertainty_network.state_dict().update(state_dict)
            # 3. load the new state dict
            self.model.uncertainty_network.load_state_dict(state_dict, strict=False)

    def freeze_layers(self):
        # Eventual fine-tuning for self-confid
        LOGGER.info("Freezing every layer except uncertainty")
        for param in self.model.named_parameters():
            if "uncertainty" in param[0]:
                print(param[0], "kept to training")
                continue
            param[1].requires_grad = False

    def disable_bn(self, verbose=False):
        # Freeze also BN running average parameters
        if verbose:
            LOGGER.info("Keeping original BN parameters")
        for layer in self.model.named_modules():
            if "bn" in layer[0] or "cbr_unit.1" in layer[0]:
                if verbose:
                    print(layer[0], "original BN setting")
                layer[1].momentum = 0
                layer[1].eval()

    def disable_dropout(self, verbose=False):
        # Freeze also BN running average parameters
        if verbose:
            LOGGER.info("Disable dropout layers to reduce stochasticity")
        for layer in self.model.named_modules():
            if "dropout" in layer[0]:
                if verbose:
                    print(layer[0], "set to eval mode")
                layer[1].eval()