# Original source code developed by @clviper
# https://github.com/clviper/droidstatx

import gzip
import hashlib
import logging
import os
import re
import time
import xml.etree.ElementTree as ET
from io import BytesIO

from elftools.elf.elffile import ELFFile

import droidsf.utils
from droidsf.smali_checks import SmaliChecks

log = logging.getLogger(__name__)

NS_ANDROID_URI = "http://schemas.android.com/apk/res/android"
NS_ANDROID = "{http://schemas.android.com/apk/res/android}"

ANDROID_CODENAMES = {
    "3": "Cupcake 1.5",
    "4": "Donut 1.6",
    "5": "Eclair 2.0",
    "6": "Eclair 2.0.1",
    "7": "Eclair 2.1",
    "8": "Froyo 2.2.x",
    "9": "Gingerbread 2.3 - 2.3.2",
    "10": "Gingerbread 2.3.3 - 2.3.7",
    "11": "Honeycomb 3.0",
    "12": "Honeycomb 3.1",
    "13": "Honeycomb 3.2.x",
    "14": "Ice Cream Sandswich 4.0.1 - 4.0.2",
    "15": "Ice Cream Sandswich 4.0.3 - 4.0.4",
    "16": "Jelly Bean 4.1.x",
    "17": "Jelly Bean 4.2.x",
    "18": "Jelly Bean 4.3.x",
    "19": "KitKat 4.4 - 4.4.4",
    "21": "Lolipop 5.0",
    "22": "Lolipop 5.1",
    "23": "Marshmallow 6.0",
    "24": "Nougat 7.0",
    "25": "Nougat 7.1",
    "26": "Oreo 8.0",
    "27": "Oreo 8.1.0",
    "28": "P ",
    "29": "Q"
}

# https://android.jlelse.eu/permission-protection-levels-969c9d0a7ebc
# https://developer.android.com/guide/topics/manifest/permission-element
# 0 - normal | 1 - dangerous | 2 - signature
ANDROID_PERMISSION_PROTECTION_LEVELS = {
    0: "Normal",
    1: "Dangerous",
    2: "Signature"
}

class DroidStatX:
    a = ""
    d = ""
    dx = ""
    xml = ""
    manifest_xml = {}
    app_xml = {}
    exportedActivities = []
    intentFilterList = {}
    componentPermissionList = {}
    activitiesWithExcludeFromRecents = []
    activitiesExtendPreferencesWithValidate = []
    activitiesExtendPreferencesWithoutValidate = []
    activitiesWithoutFlagSecure = []
    exportedReceivers = []
    exportedProviders = []
    exportedServices = []
    permissions = []
    permissionProtectionLevels = {}
    secretCodes = []
    libs = []
    assemblies = []
    assets = []
    cordova = []
    rawResources = []
    dexFiles = []
    otherFiles = []
    cordovaPlugins = []
    networkSecurityConfigDomains = []
    isAppXamarin = False
    xamarinMKBundled = False
    xamarinBundledFile = ""
    isAppCordova = False
    isAppOutsystems = False
    networkSecurityConfig = False
    minSDKVersion = ""
    targetSDKVersion = ""
    versionCode = ""
    versionName = ""
    sha256 = ""
    packageName = ""
    debuggable = False
    allowBackup = False
    certificates = []
    signatureFiles = []

    baksmaliPaths = []
    smaliChecks = None

    def __init__(self, args, apk):
        self.args = args
        self.apk = apk.apk  # androguard.core.bytecodes.apk
        self.sha256 = apk.sha256
        self.output_name = apk.output_name
        self.output_path = apk.output_path

        self.smaliChecks = SmaliChecks(self.args, apk)
        self.performAnalysis()

    def performAnalysis(self):
        log.info("Analysing manifest.xml")
        self.manifest_xml = self.apk.get_android_manifest_axml().get_xml_obj()
        self.app_xml = self.manifest_xml.findall("application")[0]

        log.info("Extracting Activities")
        self.extractActivitiesWithExcludeFromRecents()
        self.extractActivitiesWithoutSecureFlag()
        log.info("Extracting Package Properties")
        self.extractPackageProperties()
        log.info("Extracting Exported Components")
        self.extractExportedComponents()
        log.info("Extracting Permissions")
        self.extractPermissions()
        self.extractPermissionProtectionLevels()
        log.info("Extracting Certificates")
        self.extractCertificates()
        log.info("Extracting Files")
        self.extractFiles()

        self.export_analysis()

    def export_analysis(self):
        data = self.get_analysis()
        content = []
        for elem in sorted(data.items()):
            content.append(elem[0] + ": " + str(elem[1]))
        data = self.smaliChecks.get_analysis()
        for elem in sorted(data.items()):
            content.append(elem[0] + ": " + str(elem[1]))

        filename = self.output_name + ".txt"
        droidsf.utils.export_file(self.args.output_path, filename, content)
        log.info("Exported analysis to: %s", filename)

    # Return the Android Code Name for the particular Api Level.
    def getCodeName(self, apiLevel):
        return ANDROID_CODENAMES[apiLevel]

    # Extract package properties such as the minSDKVersion, PackageName, VersionName, VersionCode, isDebuggable, allowBackup
    def extractPackageProperties(self):
        usesSDK = self.manifest_xml.findall("uses-sdk")
        self.minSDKVersion = self.apk.get_min_sdk_version()
        self.targetSDKVersion = self.apk.get_target_sdk_version()
        self.packageName = self.apk.get_package()
        self.versionName = self.apk.get_androidversion_name()
        self.versionCode = self.apk.get_androidversion_code()
        if self.app_xml.get(NS_ANDROID + "debuggable") == 'true':
            self.debuggable = True
        if self.app_xml.get(NS_ANDROID + "allowBackup") == 'true':
            self.allowBackup = True
        elif self.app_xml.get(NS_ANDROID + "allowBackup") == 'false':
            self.allowBackup = False
        else:
            self.allowBackup = True

        networkSecurityConfig = self.app_xml.get(NS_ANDROID + "networkSecurityConfig")
        if networkSecurityConfig:
            self.networkSecurityConfig = True
            self.parseNetworkSecurityConfigFiles(networkSecurityConfig)

    def parseNetworkSecurityConfigFiles(self, namespace):
        xml_path = os.path.join(self.output_path, "res/xml")
        config_file = os.path.join(xml_path, "network_security_config.xml")

        xml_files = []
        if os.path.isfile(config_file):
            xml_files.append(config_file)
        elif os.path.isdir(xml_path):
            for root, dirs, files in os.walk(xml_path):
                xml_files = [os.path.join(root, f) for f in files if f.endswith(".xml")]
                # for name in files:
                #     xml_files.append(os.path.join(root, name))

        for file_path in xml_files:
            tree = ET.parse(file_path)
            root = tree.getroot()
            if self.parseNetworkSecurityConfigFile(root):
                log.info("Found network security config (%s) file: %s", namespace, file_path)

    def parseNetworkSecurityConfigFile(self, root):
        foundConfig = False

        for child in root:
            if child.tag == "base-config" or child.tag == "domain-config":
                domainConfig = {
                    'domains': [],
                    'allowClearText': True,
                    'allowUserCA': False,
                    'pinning': False,
                    'pinningExpiration': ''
                }
                if 'cleartextTrafficPermitted' in child.attrib:
                    if child.attrib['cleartextTrafficPermitted'] == "false":
                        domainConfig['allowClearText'] = False
                for sub in child:
                    if sub.tag == "domain":
                        domainConfig['domains'].append(sub.text)
                    if sub.tag == "trust-anchors":
                        for certificates in sub:
                            if certificates.attrib['src'] == "user":
                                domainConfig['allowUserCA'] = True
                    if sub.tag == "pin-set":
                        domainConfig['pinning'] = True
                        if 'expiration' in sub.attrib:
                            domainConfig['pinningExpiration'] = sub.attrib['expiration']
                self.networkSecurityConfigDomains.append(domainConfig)
                foundConfig = True

        return foundConfig

    # Create the list of permissions used by the package
    def extractPermissions(self):
        for permission in self.apk.get_permissions():
            # log.debug("Application requires permission: %s", permission)
            self.permissions.append(permission)

    def extractPermissionProtectionLevels(self):
        self.permissionProtectionLevels = {}
        for permission in self.manifest_xml.findall("permission"):
            name = permission.get(NS_ANDROID + "name")
            level = int(permission.get(NS_ANDROID + "protectionLevel"), 16)
            level_name = ANDROID_PERMISSION_PROTECTION_LEVELS[level]

            # log.debug("Found Protection Level for Permission : %s - %s", name, level_name)
            self.permissionProtectionLevels[name] = level_name

    # Extract certificates used and signature file names
    def extractCertificates(self):
        if self.apk.is_signed():
            # Test if signed v1 or v2 or both
            if self.apk.is_signed_v1() and self.apk.is_signed_v2():
                scheme_used = "v1+v2"
            elif self.apk.is_signed_v1():
                scheme_used = "v1"
            else:
                scheme_used = "v2"

        # Iterate over all certificates
        for cert in self.apk.get_certificates():
            # Each cert is now a asn1crypt.x509.Certificate object
            data = {
                "sha1": cert.sha1.hex(),  # the sha1 fingerprint
                "sha256": cert.sha256.hex(),  # the sha256 fingerprint
                "issuer": cert.issuer.human_friendly,  # issuer
                "subject": cert.subject.human_friendly,  # subject, usually the same
                "hash_algo": cert.hash_algo,  # hash algorithm
                "signature_algo": cert.signature_algo,  # Signature algorithm
                "serial_number": cert.serial_number,  # Serial number
                # "contents": cert.contents,  # The DER coded bytes of the certificate itself
            }
            self.certificates.append(data)

        # Return the name of the first signature file found.
        self.signatureFiles = self.apk.get_signature_names()

    # Create a global list of activities with the excludeFromRecents attribute
    # XXX: Checked with sweatcoin
    def extractActivitiesWithExcludeFromRecents(self):
        for activity in self.app_xml.findall("activity"):
            if activity.get(NS_ANDROID + "excludeFromRecents") == 'true':
                self.activitiesWithExcludeFromRecents.append(activity.get(NS_ANDROID + "name"))

    # Create a global list of activities that do not have the FLAG_SECURE or the excludeFromRecents attribute set.
    # XXX: Checked with pt.santandertotta.mobileparticulares
    def extractActivitiesWithoutSecureFlag(self):
        activitiesWithoutSecureFlag = []

        for activity in self.app_xml.findall("activity"):
            activityName = activity.get(NS_ANDROID + "name")

            if self.smaliChecks.doesActivityHasFlagSecure(activityName):
                log.debug("Activity %s sets FLAG_SECURE.", activityName)
            elif activityName not in self.activitiesWithExcludeFromRecents:
                self.activitiesWithoutFlagSecure.append(activityName)
            else:
                log.debug("Activity %s is on ExcludeFromRecents.", activityName)

    # Add the extracted permission of a particular component to a global list indexed by the component name.
    def extractComponentPermission(self, component):
        name = component.get(NS_ANDROID + "name")
        permission = component.get(NS_ANDROID + "permission")

        if permission:
            self.componentPermissionList[name] = permission

    # https://developer.android.com/guide/components/intents-filters
    # Create a global list with that particular object intent-filters indexed to the component name.
    def extractIntentFilters(self, component):
        name = component.get(NS_ANDROID + "name")
        items = component.findall("intent-filter")

        if not items:
            # log.debug("Component %s has no Intent Filters.", name)
            return

        filterList = []
        for item in items:
            intent_filter = {
                "action": [],
                "category": [],
                "data": []
            }
            for action in item.findall("action"):
                intent_filter["action"].append(action.get(NS_ANDROID + "name"))
            for category in item.findall("category"):
                intent_filter["category"].append(category.get(NS_ANDROID + "name"))
            for data in item.findall("data"):
                scheme = data.get(NS_ANDROID + "scheme")
                host = data.get(NS_ANDROID + "host")
                port = data.get(NS_ANDROID + "port")
                path = data.get(NS_ANDROID + "path")
                pathPattern = data.get(NS_ANDROID + "pathPattern")
                pathPrefix = data.get(NS_ANDROID + "pathPrefix")
                mimeType = data.get(NS_ANDROID + "mimeType")
                if scheme:
                    intent_filter["data"].append("scheme:" + scheme)
                    # Check presence of a SECRET_CODE
                    if scheme == "android_secret_code":
                        self.secretCodes.append(name)
                        intent_filter["data"].append("android_secret_code")
                if host:
                    intent_filter["data"].append("host:" + host)
                if port:
                    intent_filter["data"].append("port:" + port)
                if path:
                    intent_filter["data"].append("path:" + path)
                if pathPattern:
                    intent_filter["data"].append("pathPattern:" + pathPattern)
                if pathPrefix:
                    intent_filter["data"].append("pathPrefix:" + pathPrefix)
                if mimeType:
                    intent_filter["data"].append("mimeType:" + mimeType)

            filterList.append(intent_filter)

        self.intentFilterList[name] = filterList

    # Determine exported Activities taking into account the existence of exported attribute or the presence of intent-filters and check if vulnerable to Fragment Injection
    def extractExportedActivities(self):
        for activity in self.app_xml.findall("activity"):
            activityName = activity.get(NS_ANDROID + "name")

            # Extract Permission
            self.extractComponentPermission(activity)

            # Extract Intent Filters
            self.extractIntentFilters(activity)

            # Detect OutSystems applications
            if "com.outsystems.android" in activityName:
                    self.isAppOutsystems = True

            exported = activity.get(NS_ANDROID + "exported")  # 'true' / 'false' / None
            if exported == "true":
                # log.debug("Exported activity %s: found 'exported' attribute.", activityName)
                self.exportedActivities.append(activityName)
            elif exported != "false" and activityName in self.intentFilterList:
                # log.debug("Exported activity %s: found Intent Filters.", activityName)
                self.exportedActivities.append(activityName)

    def analyseExportedActivities(self):
        for activity in self.exportedActivities:
            res = self.smaliChecks.doesPreferenceActivityHasValidFragmentCheck(activity)
            if res is True:
                self.activitiesExtendPreferencesWithValidate.append(activity)
            elif res is False:
                self.activitiesExtendPreferencesWithoutValidate.append(activity)

    # Determine exported Broadcast Receivers taking into account the existence of exported attribute or the presence of intent-filters
    def extractExportedReceivers(self):
        for receiver in self.app_xml.findall("receiver"):
            receiverName = receiver.get(NS_ANDROID + "name")

            # Extract Permission
            self.extractComponentPermission(receiver)

            # Extract Intent Filters
            self.extractIntentFilters(receiver)

            exported = receiver.get(NS_ANDROID + "exported")  # 'true' / 'false' / None
            if exported == "true":
                # log.debug("Exported receiver %s: found 'exported' attribute.", receiverName)
                self.exportedReceivers.append(receiverName)
            elif exported != "false" and receiverName in self.intentFilterList:
                # log.debug("Exported activity %s: found Intent Filters.", receiverName)
                self.exportedReceivers.append(receiverName)

    # Determine exported Content Providers taking into account the existence of exported attribute or without the attributes, under API 16 they are exported by default
    def extractExportedProviders(self):
        for provider in self.app_xml.findall("provider"):
            providerName = provider.get(NS_ANDROID + "name")

            # Extract Permission
            self.extractComponentPermission(provider)

            exported = provider.get(NS_ANDROID + "exported")  # 'true' / 'false' / None
            if exported == "true":
                # log.debug("Exported provider %s: found 'exported' attribute.", providerName)
                self.exportedProviders.append(providerName)
            elif exported != "false" and self.minSDKVersion <= 16:
                # log.debug("Exported provider %s: minSDK version <= 16.", providerName)
                self.exportedProviders.append(providerName + " * In devices <= API 16 (Jelly Bean 4.1.x)")

    def analyseExportedProviders(self):
        for provider in self.exportedProviders:
            self.smaliChecks.determineContentProviderSQLi(provider)
            self.smaliChecks.determineContentProviderPathTraversal(provider)

    # Determine exported Services taking into account the existence of exported attribute or the presence of intent-filters
    def extractExportedServices(self):
        for service in self.app_xml.findall("service"):
            serviceName = service.get(NS_ANDROID + "name")

            # Extract Permission
            self.extractComponentPermission(service)

            # Extract Intent Filters
            self.extractIntentFilters(service)

            exported = service.get(NS_ANDROID + "exported")  # 'true' / 'false' / None
            if exported == "true":
                # log.debug("Exported service %s: found 'exported' attribute.", serviceName)
                self.exportedServices.append(serviceName)
            elif exported != "false" and serviceName in self.intentFilterList:
                # log.debug("Exported service %s: found Intent Filters.", serviceName)
                self.exportedServices.append(serviceName)

    # Run the functions that extract the exported components.
    def extractExportedComponents(self):
        self.extractExportedActivities()
        self.analyseExportedActivities()
        self.extractExportedReceivers()
        self.extractExportedProviders()
        self.analyseExportedProviders()
        self.extractExportedServices()

    # Determine if path is in the exclude paths configured in the config file.
    def isInExclusions(self, filename, exclusions):
        for pattern in exclusions:
            if pattern in filename:
                return True
        return False

    # Create a list of files, organized in several types and while doing it, by the existence of certain files, determine if the app is a Cordova or Xamarin app.
    def extractFiles(self):
        files = self.apk.get_files()
        try:
            for f in files:
                if self.isInExclusions(f, self.args.file_exclusions):
                    continue
                try:
                    f.encode("ascii")
                except UnicodeEncodeError as e:
                    f = f.encode('ascii', 'xmlcharrefreplace')
                if "assets/www/" in f:
                    if "assets/www/cordova.js" in f:
                        self.isAppCordova = True
                    if "assets/www/plugins/" in f:
                        beginPos = f.find("/plugins/") + 9
                        endPos = f.find("/", beginPos)
                        item = f[beginPos:endPos]
                        self.cordovaPlugins.append(item) if item not in self.cordovaPlugins else None
                    self.cordova.append(f)
                elif f[0:4] == "lib/":
                    self.libs.append(f)
                    if 'libmonodroid_bundle_app.so' in f:
                        self.isAppXamarin = True
                        self.xamarinMKBundled = True
                        self.xamarinBundledFile = f
                        self.unbundleXamarinDlls(f)
                elif f[0:11] == "assemblies/" in f:
                    self.assemblies.append(f)
                    if 'Xamarin.' in f and self.isAppXamarin is False:
                        self.isAppXamarin = True
                elif "assets/" in f:
                    self.assets.append(f)
                elif "res/raw/" in f:
                    self.rawResources.append(f)
                elif ".dex" in f:
                    self.dexFiles.append(f)
                else:
                    self.otherFiles.append(f)
        except UnicodeDecodeError as e:
            pass

    def unbundleXamarinDlls(self, lib_path):
        export_path = "xamarin_" + lib_path.split("/")[1]
        file_path = os.path.join(self.output_path, lib_path)
        output_path = os.path.join(self.output_path, export_path)
        log.info("Unpacking %s to %s.", lib_path, output_path)
        with open(file_path, "rb") as f:
            data = f.read()
            elffile = ELFFile(f)
            log.debug("Loaded ELF file: %s", lib_path)
            section = elffile.get_section_by_name('.dynsym')

            for symbol in section.iter_symbols():
                if symbol['st_shndx'] != 'SHN_UNDEF' and symbol.name.startswith('assembly_data_'):
                    start_offset = symbol['st_value']
                    end_offset = symbol['st_value'] + symbol['st_size']

                    dll_data = data[start_offset:end_offset]
                    dll_data = gzip.GzipFile(fileobj=BytesIO(dll_data)).read()

                    filename = symbol.name[14:].replace('_dll', '.dll')
                    droidsf.utils.export_file(output_path, filename, dll_data, mode="wb")
                    log.debug("Exported Xamarin DLL: %s", filename)

    # *** GETTERS ***

    def get_analysis(self):
        attrs = [
            "permissions",
            # "permissionProtectionLevels",
            "exportedActivities",
            "exportedReceivers",
            "exportedProviders",
            "exportedServices",
            "packageName",
            "minSDKVersion",
            "targetSDKVersion",
            "versionName",
            "versionCode",
            "sha256",
            "isAppCordova",
            "isAppXamarin",
            "xamarinMKBundled",
            "isAppOutsystems",
            "debuggable",
            "allowBackup",
            "networkSecurityConfig",
            "isMultiDex",
            "assets",
            "cordova",
            "rawResources",
            "libs",
            "otherFiles",
            "assemblies",
            "dexFiles",
            "secretCodes",
            "intentFilterList",
            "networkSecurityConfigDomains",
            "activitiesWithExcludeFromRecents",
            "activitiesWithoutFlagSecure",
            "activitiesExtendPreferencesWithValidate",
            "activitiesExtendPreferencesWithoutValidate",
            "cordovaPlugins",
            "certificates",
            "signatureFiles"
        ]
        data = {}
        for attr in attrs:
            value = getattr(self, attr)
            if isinstance(value, list):
                try:
                    value = sorted(value)
                except TypeError:
                    pass
            elif callable(value):
                value = value()

            data[attr] = value

        return data

    # Return the permission defined in the particular component.
    # XXX: unused
    def getComponentPermission(self, name):
        try:
            return self.componentPermissionList[name]
        except:
            return ""

    # Return the app permissions global list
    def getPermissions(self):
        return self.permissions

    # Return the exported activities global list
    def getExportedActivities(self):
        return self.exportedActivities

    # Return the the exported broadcast receivers global list
    def getExportedReceivers(self):
        return self.exportedReceivers

    # Return the exported content providers global list
    def getExportedProviders(self):
        return self.exportedProviders

    # Return the exported services global list
    def getExportedServices(self):
        return self.exportedServices

    # Return the app package name
    def getPackageName(self):
        return self.packageName

    # Return the app minSDKVersion
    def getMinSDKVersion(self):
        return self.minSDKVersion

    # Return the app targetSDKVersion
    def getTargetSDKVersion(self):
        return self.targetSDKVersion

    # Return the app versionName
    def getVersionName(self):
        return self.versionName

    # Return the app versionCode
    def getVersionCode(self):
        return self.versionCode

    # Return the APK SHA256
    def getSHA256(self):
        return self.sha256

    def isCordova(self):
        if self.isAppCordova is True:
            return "Yes"
        else:
            return "No"

    def isXamarin(self):
        if self.isAppXamarin is True:
            return "Yes"
        else:
            return "No"

    def isXamarinBundled(self):
        if self.xamarinMKBundled is True:
            return "Yes"
        else:
            return "No"

    def isOutsystems(self):
        if self.isAppOutsystems is True:
            return "Yes"
        else:
            return "No"

    def isDebuggable(self):
        if self.debuggable is True:
            return "Yes"
        else:
            return "No"

    def isBackupEnabled(self):
        if self.allowBackup is True:
            return "Yes"
        else:
            return "No"

    def hasNetworkSecurityConfig(self):
        return self.networkSecurityConfig

    def isMultiDex(self):
        return len(self.dexFiles) > 1

    def getAssets(self):
        return self.assets

    def getCordovaFiles(self):
        return self.cordova

    def getRawResources(self):
        return self.rawResources

    def getLibs(self):
        return self.libs

    def getOtherFiles(self):
        return self.otherFiles

    def getXamarinAssemblies(self):
        return self.assemblies

    def getDexFiles(self):
        return self.dexFiles

    def getSecretCodes(self):
        return self.secretCodes

    def getIntentFiltersList(self):
        return self.intentFilterList

    def getNetworkSecurityConfigDomains(self):
        return self.networkSecurityConfigDomains

    def getActivitiesWithExcludeFromRecents(self):
        return self.activitiesWithExcludeFromRecents

    def getActivitiesWithoutSecureFlag(self):
        return self.activitiesWithoutFlagSecure

    def getActivitiesExtendPreferencesWithValidate(self):
        return self.activitiesExtendPreferencesWithValidate

    def getActivitiesExtendPreferencesWithoutValidate(self):
        return self.activitiesExtendPreferencesWithoutValidate

    def getCordovaPlugins(self):
        return self.cordovaPlugins


import tkinter as tk
from graph import Graph

class MetroTravelUI:
    def __init__(self, graph, airport_codes):
        # Initialize the UI with a graph and a list of airport codes
        self.graph = graph
        self.airport_codes = airport_codes
        self.root = tk.Tk()
        self.root.title("Metro Travel")

        # Create UI elements for origin and destination inputs
        self.origin_label = tk.Label(self.root, text="Origen:")
        self.origin_label.grid(row=0, column=0)
        self.origin_entry = tk.Entry(self.root)
        self.origin_entry.grid(row=0, column=1)

        self.destination_label = tk.Label(self.root, text="Destino:")
        self.destination_label.grid(row=1, column=0)
        self.destination_entry = tk.Entry(self.root)
        self.destination_entry.grid(row=1, column=1)

        # Create a UI element for the visa checkbox
        self.visa_var = tk.BooleanVar()
        self.visa_check = tk.Checkbutton(self.root, text="Tiene visa", variable=self.visa_var)
        self.visa_check.grid(row=2, column=0, columnspan=2)

        # Create a UI element for the result label
        self.result_label = tk.Label(self.root, text="")
        self.result_label.grid(row=3, column=0, columnspan=2)

        # Create a UI element for the calculate button
        self.calculate_button = tk.Button(self.root, text="Calcular ruta", command=self.calculate_route)
        self.calculate_button.grid(row=4, column=0, columnspan=2)

    def run(self):
        # Start the UI main loop
        self.root.mainloop()

    def calculate_route(self):
        # Get the origin, destination, and visa checkbox values from the UI
        origin = self.origin_entry.get().strip().upper()
        destination = self.destination_entry.get().strip().upper()
        has_visa = self.visa_var.get()

        # Check if the origin and destination codes are valid
        if origin not in self.airport_codes or destination not in self.airport_codes:
            self.result_label.config(text="Origen o destino no válido.")
            return

        # Find the indices of the origin and destination codes in the airport codes list
        src_idx = self.airport_codes.index(origin)
        dest_idx = self.airport_codes.index(destination)

        # Find the cheapest path between the origin and destination airports
        cost, path = self.graph.find_route(src_idx, dest_idx, has_visa)

        # Update the result label with the path and cost
        if cost == float('inf'):
            self.result_label.config(text="No se encontró una ruta válida.")
        else:
            path = [self.airport_codes[i] for i in path]
            self.result_label.config(text=f"Ruta: {' -> '.join(path)}\nCosto: ${cost}")


def inclusionExclusion3(S, T, U): #creates inclusion-exclusion function with lists S, T, and U as inputs
    cardinality = 0 #sets the cardinality of the union of sets to 0
    allSets = intersect(U, intersect(S, T)) #creates a list that is the intersection of the inputs
    for s in S: #loops through each element in S
        cardinality += 1 #adds 1 to the cardinality for each element in S
    for t in T: # loops through each element in T
        cardinality += 1 #adds 1 to the cardinality for each element in T
    for u in U: #loops through each element in U
        cardinality += 1 #adds 1 for each element in U
    for element in intersect(S, T): #loops through each element in the intersection of S and T
        cardinality -=1 #subtracts 1 from the cardinality for each element in S and T
    for element in intersect(S, U): #loops through each element in the intersection of S and U
        cardinality -=1 #subtracts 1 from the cardinality for each element in S and U
    for element in intersect(U, T):#loops through each element in the intersection of U and T
        cardinality -= 1 #subtracts 1 from the cardinality for each element in U and T
    for element in allSets: #loops through each element in the intersection of each set
        cardinality += 1 #adds 1 to the cardinality for each element in the intersection of each set
    return cardinality #returns the cardinality

def additive3(S, T, U): #creates the additive3 function with inputs S, T, and U
    if(len(intersect(S, T)) != 0): #creates a condition for if the cardinality of the intersection of S and T does not equal 0
        return False #returns false if the condition holds true
    if(len(intersect(S, U)) != 0): #creates a condition for if the cardinality of the intersection of S and U does not equal 0
        return False #returns false if the condition holds true
    if(len(intersect(T, U)) != 0): #creates a condition for if the cardinality of the intersection of U and T does not equal 0
        return False #returns false if the condition holds true
    return len(S) + len(T) + len(U) #returns the cardinality of each set

#below is the intersection function from coding assignment 2 that will help the functions above
def intersect(A, B): #creates intersect funtion with lists A & B as arguments
    temp_list = [] #creates temporary list to be used in the function
    intersect_list = [] #creates list that will be intersection of A & B
    for element in A: #goes through elements of A
        if element in B: #creates a condition that makes it so element must also be in B
            temp_list.append(element) #adds elements to temporary list
    for element in temp_list: #goes through elements in the temporary list
        if element not in intersect_list: #creates a condition for anything not in list already
            intersect_list.append(element) #adds elements to final list
            intersect_list.sort() #sorts list in ascending order
    return intersect_list #returns the final list


import parameter as par
import util
from c2x import C2X
from direction import Direction
from key import Key
from message import Message
from state import State

class PlatoonVehicle:
    """
    Class that handles platoon members and followers.
    """

    def __init__(self, v_id: str, c2x: C2X):
        self.v_id = v_id
        self.step = 0
        self.timer = 0
        self.state = State.NO_STATE
        self.direction = Direction.LEFT
        self.in_lane_change = False
        self.begin_lane_change = False
        self.in_change_back = False
        self.begin_change_back = False
        self.msg_queue = []
        self.rec_data = {Key.STATES_VISITED: [],
                         Key.OT_TIME: [],
                         Key.LANES_VISITED: [],
                         Key.SPEED_MIN: float('inf'),
                         Key.SPEED_MAX: float('-inf'),
                         Key.MIN_GAP: float('inf')}
        # Messaging system
        self.c2x = c2x
        # Reset timer and overtaking state
        self.set_timer(0)
        self.set_state(State.IDLE)
        # Lane change from lane lc_from to lane lc_to
        self.lc_from = 0
        self.lc_to = 0
        # Platoon did cross the lane markings -> self.in_new_lane = True
        self.in_new_lane = False
        self.lane_offset_last_step = 0
        self.abort_message = ''

    def detect_lane_crossing(self):
        """
        Detects whether the platoon crosses the lane markings during a lane change maneuver.
        """
        # Detect crossing of lane markings
        curr_offset = util.get_lateral_offset(self.v_id)
        # Lane crossed
        if abs(curr_offset) > par.LANE_OFFSET_MIN and self.lane_offset_last_step * curr_offset < 0:
            self.in_new_lane = True
        self.lane_offset_last_step = curr_offset

    def safety_dist_platoon(self):
        """
        Computes the necessary safety distance (in meters) of the platoon to a (non-platooning) vehicle in front.

        :return: Necessary safety distance in meters
        """
        d_safety = util.get_tau(self.v_id) * util.get_speed(self.v_id)
        # Raise d_safety a little bit when not in a lane change to avoid oscillation
        if not self.in_lane_change:
            d_safety *= (1 + par.OSCILLATION_MOD)
        return d_safety

    def alpha_area_free(self, debug=True) -> bool:
        """
        Returns True if the FL (or FR, depending on self.direction) area is considered free, False otherwise. The
        area is considered free if no vehicle is detected in that area or the closest vehicle in that area is at
        least the platoon's safety distance away.

        :return: True if FL (or FR, depending on self.direction) area is considered free, False otherwise
        """
        is_free = True
        # Get relevant vehicle in FL (or FR) area, platoon is changing lanes and already in new lane
        if self.in_lane_change and self.in_new_lane:
            neighbor = util.get_leader(self.v_id)
        else:
            neighbor = util.get_neighbor(self.v_id, self.direction, front=True)
        if neighbor is not None:
            if neighbor[1] < self.safety_dist_platoon():
                is_free = False
                self.debug_print(f'F{self.direction.short()} area is occupied by {neighbor[0]} with distance '
                                 f'{round(neighbor[1], 2)} < safety={round(self.safety_dist_platoon(), 2)}')
            else:
                if debug:
                    self.debug_print(f'F{self.direction.short()} area is free ({neighbor[0]} with distance '
                                     f'{round(neighbor[1], 2)} >= '
                                     f'safety={round(self.safety_dist_platoon(), 2)})')
        else:
            if debug:
                self.debug_print(f'F{self.direction.short()} area is free (no vehicle detected)')
        return is_free

    def d_min(self, a: float, follower) -> float:
        """
        Returns the minimum distance the vehicle in the RL (or RR) area has to be away from the platoon to ensure
        that this vehicle only has to break with a (negative) acceleration of at most <a> and that the safety distance
        is observed.

        :param a: Minimum allowed (negative) acceleration of the RL (or RR) vehicle
        :param follower: Tuple of follower's id and distance
        :return: Minimum distance to the RL (or RR) vehicle so the beta area is considered free
        """
        v_p = util.get_speed(self.v_id)
        v_r = util.get_speed(follower[0])
        T_r = par.R_REACT
        T_g = par.R_TIME_GAP
        if v_r > v_p and a < 0:
            result = (-1 / (2 * a)) * (v_p - v_r) ** 2 + v_r * T_r + v_p * T_g
        elif v_r <= v_p and a <= 0:
            result = v_r * (T_r + T_g)
        else:
            result = float('inf')
        # Take truck min gap of 50m into account when changing to the right
        if self.direction == Direction.RIGHT:
            result = max(result, par.TRUCK_MIN_GAP)
        # Raise d_min a little bit when not in a lane change to avoid oscillation
        if not self.in_lane_change:
            result *= (1 + par.OSCILLATION_MOD)
        return result

    def beta_area_free(self, debug=True) -> bool:
        """
        Returns True if the RL (or RR, depending on self.direction) area is considered free, False otherwise.
        The area is considered free if no vehicle is detected in that area or the closest vehicle is at least d_min
        away.

        :return: True if RL (or RR, depending on self.direction) area is considered free, False otherwise
        """
        is_free = True
        # Max accepted (negative) acceleration by neighbor follower
        if self.in_lane_change:
            threshold = par.MAX_ACCEL_R_VEHICLE_ABORT
        else:
            if self.direction == Direction.LEFT:
                threshold = par.MAX_ACCEL_VEHICLE_LEFT
            else:
                threshold = par.MAX_ACCEL_VEHICLE_RIGHT
        # Get relevant vehicle in RL (or RR) area
        if self.in_lane_change and self.in_new_lane:
            # Platoon in lane change and already in new lane
            neighbor = util.get_follower(self.v_id)
        else:
            # Platoon not in lange change or in lane change and still in old lane
            neighbor = util.get_neighbor(self.v_id, self.direction, front=False)
        if neighbor is not None:
            # Compute necessary distance to closest vehicle in RL (or RR) area to ensure safety distance
            min_dist = self.d_min(threshold, neighbor)
            if debug:
                self.debug_print(f"R{self.direction.short()} '{neighbor[0]}' needs distance "
                                 f'{round(min_dist, 2)} (has {round(neighbor[1], 2)}) '
                                 f'to ensure safety distance (threshold={threshold})')
            if neighbor[1] < min_dist:
                self.debug_print(f'R{self.direction.short()} area is occupied by {neighbor[0]} with distance '
                                 f'{round(neighbor[1], 2)} (needs {round(min_dist, 2)})')
                is_free = False
            else:
                if debug:
                    self.debug_print(f'R{self.direction.short()} area is free ({neighbor[0]})')
        else:
            if debug:
                self.debug_print(f'R{self.direction.short()} area is free (no vehicle detected)')
        return is_free

    def exec_lane_change(self):
        """
        Orders the vehicle to changes lanes in the direction of self.direction.
        """
        # Vehicle is already in target lane
        if self.in_new_lane:
            lat_dist = -1 * util.get_lateral_offset(self.v_id)
        # Vehicle is still in original lane
        else:
            if self.direction == Direction.LEFT:
                lat_dist = par.LANE_WIDTH - util.get_lateral_offset(self.v_id)
            else:
                lat_dist = -1 * (par.LANE_WIDTH + util.get_lateral_offset(self.v_id))
        util.change_sublane(self.v_id, lat_dist)
        self.debug_print(f'Lane change from lane {self.lc_from} to {self.lc_to} started (lat={round(lat_dist, 2)})')

    def set_state(self, state: State):
        """
        Sets the lane change state of the vehicle according to the cooperative overtaking algorithm.
        """
        self.state = state
        # Distinguish between left and right lane change maneuver states
        # (only for testing the algorithm, not needed for the actual algorithm)
        if self.direction == Direction.LEFT or state == State.PASSING or state == State.OVERTAKING_COMPLETE:
            self.rec_data[Key.STATES_VISITED].append(state)
        else:
            self.rec_data[Key.STATES_VISITED].append(State['R_' + state.name])
        self.debug_print(f'*** Switching to state {state.value}: {state} ***')

    def set_timer(self, value=par.MAX_TIMER):
        """
        Sets the timer to its standard value. If no value is given, par.MAX_TIMER is used.
        """
        self.timer = value
        self.debug_print(f'Timer set to {self.timer}')

    def decrement_timer(self):
        """
        Reduces the timer of each vehicle by 1 until it reaches 0.
        """
        if self.timer > 0:
            self.timer -= 1

    def record(self):
        """
        Records data of the platoon vehicle.
        """
        # Current lane
        lane = util.get_lane_index(self.v_id)
        if len(self.rec_data[Key.LANES_VISITED]) == 0 or self.rec_data[Key.LANES_VISITED][-1] != lane:
            self.rec_data[Key.LANES_VISITED].append(lane)
        # Min and max speed
        speed = round(util.get_speed(self.v_id), 2)
        if self.rec_data[Key.SPEED_MIN] > speed:
            self.rec_data[Key.SPEED_MIN] = speed
        if self.rec_data[Key.SPEED_MAX] < speed:
            self.rec_data[Key.SPEED_MAX] = speed

    def front_or_rear_member_not_following(self):
        """
        Returns True if the platoon member in front or in rear is not following laterally anymore
        (meaning the lateral offset is bigger than par.MEMBER_MAX_OFFSET).

        :return: True if front or rear member is not following laterally
        """
        # Own lateral position diverges too much from the leader's lateral position
        own_index = int(self.v_id.split('_')[1])
        own_lat_pos = abs(util.get_lateral_offset(self.v_id))
        if own_index > 0:
            front_id = par.ID_PRE + '_' + str(own_index - 1)
            front_lat_pos = abs(util.get_lateral_offset(front_id))
            if abs(own_lat_pos - front_lat_pos) > par.MEMBER_MAX_OFFSET:
                return True
        if own_index < par.PLATOON_MEMBERS - 1:
            rear_id = par.ID_PRE + '_' + str(own_index + 1)
            rear_lat_pos = abs(util.get_lateral_offset(rear_id))
            if abs(own_lat_pos - rear_lat_pos) > par.MEMBER_MAX_OFFSET:
                return True
        return False

    def next_step(self):
        """
        Main next steps that are executed for both the leader and the followers.
        """
        self.step += 1
        # Detect crossing of a lane marking
        if self.in_lane_change or self.in_change_back:
            self.detect_lane_crossing()
        self.record()
        # Discard earlier received messages
        self.msg_queue.clear()
        self.receive_messages()
        self.decrement_timer()

    # *** Messaging ***

    def receive_messages(self):
        """
        Receives all relevant messages and adds them to the queue.
         """
        self.msg_queue.extend(self.c2x.receive(self.v_id, self.step))
        for msg in self.msg_queue:
            self.debug_print(f'Message in queue from {msg["src"]}: {msg["msg"].name}')

    def get_messages_by_type(self, msg_type: Message):
        """
        Returns all messages from the queue that are of type msg. Deletes these messages from the queue.
        """
        result = [msg for msg in self.msg_queue if msg['msg'] == msg_type]
        self.msg_queue = [msg for msg in self.msg_queue if msg['msg'] != msg_type]
        return result

    def send_message(self, dest, msg: Message, data=None):
        """
        Sends a message to dest via the C2X communication system.
        """
        self.c2x.send(self.v_id, dest, msg, data, self.step)
        self.debug_print(f'Message sent to {dest}: \'{msg.name}\'')

    # *** Debug methods ***

    def print(self, substates=False):
        """
        Print some information about the vehicle (for debugging purposes).
        """
        print(f'  {self.v_id}: state sequence: {self.rec_data[Key.STATES_VISITED]}')

    def debug_print(self, message):
        """
        Prints debug information to stdout.
        """
        if par.DEBUG:
            print(f'({self.step}) {self.v_id}: {message}')

    # *** States that apply to both the leader and the follower ***

    def changing_lane(self):
        """
        State "changing lane"
        """
        # Begin lane change
        if self.begin_lane_change:
            self.begin_lane_change = False
            self.in_lane_change = True
            self.in_new_lane = False
            self.lc_from = util.get_lane_index(self.v_id)
            self.lc_to = self.lc_from + self.direction.value
            self.exec_lane_change()
        # During lane change
        elif self.in_lane_change:
            # Handle disruptions during lane change
            if util.get_lateral_speed(self.v_id) != 0:
                # FL (or FR) area occupied
                if not self.alpha_area_free(debug=False):
                    self.begin_change_back = True
                    self.abort_message = 'Front area occupied'
                # RL (or RR) area occupied
                if not self.beta_area_free(debug=False):
                    self.begin_change_back = True
                    self.abort_message = 'Rear area occupied'
                # Unreliable messaging: another platoon member does not change lanes
                if self.front_or_rear_member_not_following():
                    self.begin_change_back = True
                    self.abort_message = 'A member is not following laterally anymore'
            # Lane change done
            elif self.in_new_lane and util.is_centered(self.v_id):
                self.lc_from = self.lc_to
                self.in_lane_change = False

    def lane_change_complete(self):
        """
        State: "lane change complete"
        """
        self.in_new_lane = False
        self.direction = self.direction.flip()

    def changing_back(self):
        """
        State: "changing back"
        """
        # Begin changing back
        if self.begin_change_back:
            self.begin_change_back = False
            self.in_change_back = True
            self.in_lane_change = False
            self.direction = self.direction.flip()
            self.lc_from = util.get_lane_index(self.v_id)
            # Platoon already in new lane
            if self.in_new_lane:
                self.lc_to = self.lc_from + self.direction.value
            else:
                self.lc_to = self.lc_from
            # When already in new lane -> now in old lane
            # When still in old lane -> now in new lane
            self.in_new_lane = not self.in_new_lane
            self.exec_lane_change()
        # During changing back
        elif self.in_change_back:
            # Lane change done
            if self.in_new_lane and util.is_centered(self.v_id):
                self.direction = self.direction.flip()
                self.in_change_back = False


import argparse
import copy
import abc
import numpy as np
import time
import sys

import openmesh as om
from openmesh import TriMesh

import scipy.sparse as sp
from scipy.sparse import coo_matrix
from scipy.sparse.linalg import eigs

from sksparse.cholmod import cholesky

class Solver:
    """
    Base class of Linear System solver
    """
    def __init__(self, args):
        self.save_file = args.output
        self.input_file = args.input
        self.log = args.verbose

        self.lamba = args.l
        self.kappa = args.k

        self.beta_max = args.beta_max

        self.beta0 = args.beta
        self.beta = self.beta0

        self.input_mesh = om.read_trimesh(self.input_file)

        self.mesh = None
        self._type = None
        self.type = None
        self.t_solver = 'ch'

    def prepare_mesh(self, mesh=None):
        '''
        self.mesh is mesh to be processed
        '''
        if mesh==None:
            self.mesh = copy.deepcopy(self.input_mesh)
        else:
            self.mesh = copy.deepcopy(mesh)

        # set automatic parameters
        if self.lamba == None:
            avl = average_lenth(self.mesh)
            self.lamba = 0.02 * avl**2

    def solve(self):
        print("[ {} ]".format(self._type))
        start_time = time.time()

        print("Processing mesh with ")
        print("vertices: ", self.mesh.n_vertices())
        print("edges: ", self.mesh.n_edges())
        print("faces: ", self.mesh.n_faces())

        self.optimize()

        final_time = time.time()
        self.duration = final_time - start_time
        print("Optimize Time: %f (s)" % (self.optimize_time))
        print("Total Time: %f (s)" % (self.duration))
        print("Iterations: %d" % (self.iteration))

    def save(self, mesh, path=None):
        if path==None:
            om.write_mesh(self.save_file, mesh)
        else:
            om.write_mesh(path, mesh)

    @abc.abstractmethod
    def build_operator(self):
        """
        this paper contain three types of operator:
         - vertex-based cotangent operator
         - cotangent edge operator
         - area-based edge operator

        return: D, p, p*
        """

    def optimize(self) -> TriMesh:
        mesh = self.mesh
        N_edge = mesh.n_edges()
        N_vert = mesh.n_vertices()

        if self.type == 'edge':
            N = N_edge
        else: # type=='vert
            N = N_vert

        D, p, p_star = self.build_operator()
        start_time = time.time()
        # --------------build linear system--------------
        cnt = 1
        while self.beta < self.beta_max:
            # stage 1:
            delta = D @ p

            for i in range(N):
                if delta[i][0]**2 + delta[i][1]**2 + delta[i][2]**2 < self.lamba / self.beta:
                    delta[i][0] = delta[i][1] = delta[i][2] = 0.0

            # stage 2:
            A = sp.eye(N_vert) + self.beta * (D.transpose() @ D)
            b = p_star + self.beta * (D.transpose() @ delta)

            if self.t_solver == 'ch':
                factor = cholesky(A)
            # p, info = sp.linalg.cg(A, b, x0=p)
            for i in range(3):
                if self.t_solver == 'cg':
                    p[:,i], info = sp.linalg.cg(A, b[:,i], x0=p[:,i])
                if self.t_solver == 'ch':
                    # scipy cholesky
                    # y = scipy.linalg.solve_triangular(L, b[:,i], lower=True, check_finite=False)
                    # p[:,i] = scipy.linalg.solve_triangular(L.T, y, lower=True, check_finite=False)
                    p[:,i] = factor.solve_A(b[:,i])
            if self.log:
                print('iter {}: beta={}, p[0,0]={}'.format(cnt, self.beta, p[0][0]))
            cnt += 1
            self.beta *= self.kappa
        self.denoised_vert = p
        final_time = time.time()
        self.optimize_time = final_time - start_time
        self.iteration = cnt
        # --------------return mesh--------------
        for i, vh in enumerate(mesh.vertices()):
            mesh.set_point(vh, p[i])
        self.denoised_mesh = mesh
        return mesh

#---------------------utils--------------------------#

def average_lenth(mesh) -> float:
        lenth = 0
        for eh in mesh.edges():
            lenth += mesh.calc_edge_length(eh)
        return lenth / mesh.n_edges()

def add_noise(mesh, l=0.3, set_seed=777):
    """
    return a new mesh with noise of l * average edge length
    """
    np.random.seed(set_seed)
    input_mesh_avglen = average_lenth(mesh)
    ret = copy.deepcopy(mesh)
    for vh in ret.vertices():
        point = ret.point(vh)
        noise = np.random.normal(0, input_mesh_avglen * l)
        new_point = point + noise
        ret.set_point(vh, new_point)
    return ret

#---------------------Methods--------------------------#

class VertexSolver(Solver):
    def __init__(self, args):
        super(VertexSolver, self).__init__(args)
        self._type = 'Vertex Based Method'
        self.type = 'vert'

    def build_operator(self):
        mesh = self.mesh
        N_edge = mesh.n_edges()
        N_vert = mesh.n_vertices()

        # --------------build D--------------
        D_row = []
        D_col = []
        D_data = []

        for cnt, vh in enumerate(mesh.vertices()):
            neighbor_vhs = [vh for vh in mesh.vv(vh)]
            cots = []
            for neighbor_vh in neighbor_vhs:
                heh = mesh.find_halfedge(vh, neighbor_vh)
                vh1 = mesh.to_vertex_handle(heh)
                vh3 = mesh.from_vertex_handle(heh)
                vh2 = mesh.opposite_vh(heh)
                vh4 = mesh.opposite_he_opposite_vh(heh)

                p1, p2, p3, p4 = mesh.point(vh1), mesh.point(vh2), mesh.point(vh3), mesh.point(vh4)

                cot123 = np.dot(p1-p2, p3-p2)/np.linalg.norm(np.cross(p1-p2, p3-p2))
                cot341 = np.dot(p3-p4, p1-p4)/np.linalg.norm(np.cross(p3-p4, p1-p4))

                cots.append(cot123+cot341)
            cot_sum = sum(cots)
            for i, neighbor_vh in enumerate(neighbor_vhs):
                D_row.append(cnt)
                D_col.append(neighbor_vh.idx())
                D_data.append(cots[i]/cot_sum)
            D_row.append(cnt)
            D_col.append(vh.idx())
            D_data.append(-1)

        D = coo_matrix((D_data, (D_row, D_col)), shape=(N_vert, N_vert))
        p = np.array([mesh.point(vertex) for vertex in mesh.vertices()])
        p_star = copy.deepcopy(p)
        return D, p, p_star

class VertexSolver_FFT(VertexSolver):
    def __init__(self, args):
        super(VertexSolver_FFT, self).__init__(args)
        self._type = 'Vertex Based Method (Using FFT)'

    def decomposite(self, L):
        print('computing eigen decomposition')
        num_eigen = 100
        eignvalue, V = eigs(L, k=num_eigen)
        E = np.diag(eignvalue)
        print("done decomposition")
        return V, E, V.transpose()

    def optimize(self) -> TriMesh:
        mesh = self.mesh
        N_edge = mesh.n_edges()
        N_vert = mesh.n_vertices()

        if self.type == 'edge':
            N = N_edge
        else: # type=='vert
            N = N_vert

        D, p, p_star = self.build_operator()
        U, _, UT = self.decomposite(D)

        MTF = np.power(np.linalg.norm(UT@(D.transpose()@D), ord='fro'), 2)
        start_time = time.time()
        # --------------build linear system--------------
        cnt = 1
        while self.beta < self.beta_max:
            # stage 1:
            delta = D @ p

            for i in range(N):
                if delta[i][0]**2 + delta[i][1]**2 + delta[i][2]**2 < self.lamba / self.beta:
                    delta[i][0] = delta[i][1] = delta[i][2] = 0.0

            # stage 2:
            div = 1 + self.beta * MTF
            for i in range(3):
                res = UT@p_star[:,i] + self.beta * UT @ (D.transpose() @ delta[:,i])
                res = res / div
                res = U @ res
                p[:,i] = res

            if self.log:
                print('iter {}: beta={}, p[0,0]={}'.format(cnt, self.beta, p[0][0]))
            cnt += 1
            self.beta *= self.kappa
        self.denoised_vert = p
        final_time = time.time()
        self.optimize_time = final_time - start_time
        self.iteration = cnt
        # --------------return mesh--------------
        for i, vh in enumerate(mesh.vertices()):
            mesh.set_point(vh, p[i])
        self.denoised_mesh = mesh
        return mesh

class CotEdgeSolver(Solver):
    def __init__(self, args):
        super(CotEdgeSolver, self).__init__(args)
        self._type = 'Edge Based Method'
        self.type = 'edge'

    def build_operator(self):
        mesh = self.mesh
        N_edge = mesh.n_edges()
        N_vert = mesh.n_vertices()

        # --------------build D--------------
        D_row = []
        D_col = []
        D_data = []

        for cnt, eh in enumerate(mesh.edges()):
            heh1 = mesh.next_halfedge_handle(mesh.halfedge_handle(eh, 0)) # next edge of 1st halfedge of e
            heh2 = mesh.next_halfedge_handle(mesh.halfedge_handle(eh, 1)) # next edge of 2nd halfedge of e

            vh3, vh4 = mesh.from_vertex_handle(heh1), mesh.to_vertex_handle(heh1)
            vh1, vh2 = mesh.from_vertex_handle(heh2), mesh.to_vertex_handle(heh2)
            # corresponding to Figure 3
            p1, p2, p3, p4 = mesh.point(vh1), mesh.point(vh2), mesh.point(vh3), mesh.point(vh4)

            S123 = 0.5 * np.linalg.norm(np.cross(p2 - p1, p2 - p3))
            S134 = 0.5 * np.linalg.norm(np.cross(p4 - p1, p4 - p3))
            l13 = np.linalg.norm(p1 - p3)

            cot231 = np.dot(p2-p3, p1-p3)/np.linalg.norm(np.cross(p2-p3, p1-p3))
            cot134 = np.dot(p1-p3, p4-p3)/np.linalg.norm(np.cross(p1-p3, p4-p3))
            cot312 = np.dot(p3-p1, p2-p1)/np.linalg.norm(np.cross(p3-p1, p2-p1))
            cot413 = np.dot(p4-p1, p3-p1)/np.linalg.norm(np.cross(p4-p1, p3-p1))

            coef1 = -1 * cot231 - cot134
            coef2 = cot231 + cot312
            coef3 = -1 * cot312 - cot413
            coef4 = cot134 + cot413

            D_row.append(cnt)
            D_col.append(vh1.idx())
            D_data.append(coef1)

            D_row.append(cnt)
            D_col.append(vh2.idx())
            D_data.append(coef2)

            D_row.append(cnt)
            D_col.append(vh3.idx())
            D_data.append(coef3)

            D_row.append(cnt)
            D_col.append(vh4.idx())
            D_data.append(coef4)

        D = coo_matrix((D_data, (D_row, D_col)), shape=(N_edge, N_vert))
        p = np.array([mesh.point(vertex) for vertex in mesh.vertices()])
        p_star = copy.deepcopy(p)
        return D, p, p_star

class AreaEdgeSolver(Solver):
    def __init__(self, args):
        super(AreaEdgeSolver, self).__init__(args)
        self._type = 'Area Based Method'
        self.type = 'edge'

    def build_operator(self):
        mesh = self.mesh
        N_edge = mesh.n_edges()
        N_vert = mesh.n_vertices()

        # --------------build D--------------
        D_row = []
        D_col = []
        D_data = []

        for cnt, eh in enumerate(mesh.edges()):
            heh1 = mesh.next_halfedge_handle(mesh.halfedge_handle(eh, 0)) # next edge of 1st halfedge of e
            heh2 = mesh.next_halfedge_handle(mesh.halfedge_handle(eh, 1)) # next edge of 2nd halfedge of e

            vh3, vh4 = mesh.from_vertex_handle(heh1), mesh.to_vertex_handle(heh1)
            vh1, vh2 = mesh.from_vertex_handle(heh2), mesh.to_vertex_handle(heh2)
            # corresponding to Figure 3
            p1, p2, p3, p4 = mesh.point(vh1), mesh.point(vh2), mesh.point(vh3), mesh.point(vh4)

            S123 = 0.5 * np.linalg.norm(np.cross(p2 - p1, p2 - p3))
            S134 = 0.5 * np.linalg.norm(np.cross(p4 - p1, p4 - p3))
            l13 = np.linalg.norm(p1 - p3)

            coef1 = (S123 * np.dot(p4 - p3, p3 - p1) + S134 * np.dot(p1 - p3, p3 - p2)) / (l13 * l13 * (S123 + S134))
            coef2 = S134 / (S123 + S134)
            coef3 = (S123 * np.dot(p3 - p1, p1 - p4) + S134 * np.dot(p2 - p1, p1 - p3)) / (l13 * l13 * (S123 + S134))
            coef4 = S123 / (S123 + S134)

            D_row.append(cnt)
            D_col.append(vh1.idx())
            D_data.append(coef1)

            D_row.append(cnt)
            D_col.append(vh2.idx())
            D_data.append(coef2)

            D_row.append(cnt)
            D_col.append(vh3.idx())
            D_data.append(coef3)

            D_row.append(cnt)
            D_col.append(vh4.idx())
            D_data.append(coef4)

        D = coo_matrix((D_data, (D_row, D_col)), shape=(N_edge, N_vert))
        p = np.array([mesh.point(vertex) for vertex in mesh.vertices()])
        p_star = copy.deepcopy(p)
        return D, p, p_star

class AreaEdgeSolver_R(Solver):
    def __init__(self, args):
        super(AreaEdgeSolver_R, self).__init__(args)
        self._type = 'Area Based Method w/ reg'
        self.type = 'edge'
        self.alpha = args.a

    def get_gemma(self, mesh):
        PI = 3.14159265359
        cnt, sum = 0, 0
        for heh in mesh.halfedges():
            angle = mesh.calc_dihedral_angle(heh) * (180/PI)
            sum += abs(angle)
            cnt += 1
        return sum/cnt

    def build_operator(self):
        mesh = self.mesh
        N_edge = mesh.n_edges()
        N_vert = mesh.n_vertices()

        if self.alpha == 0:
            self.alpha = 0.1 * self.get_gemma(mesh)

        # --------------build D and R--------------
        D_row = []
        D_col = []
        D_data = []

        R_data = []

        for cnt, eh in enumerate(mesh.edges()):
            heh1 = mesh.next_halfedge_handle(mesh.halfedge_handle(eh, 0)) # next edge of 1st halfedge of e
            heh2 = mesh.next_halfedge_handle(mesh.halfedge_handle(eh, 1)) # next edge of 2nd halfedge of e

            vh3, vh4 = mesh.from_vertex_handle(heh1), mesh.to_vertex_handle(heh1)
            vh1, vh2 = mesh.from_vertex_handle(heh2), mesh.to_vertex_handle(heh2)
            # corresponding to Figure 3
            p1, p2, p3, p4 = mesh.point(vh1), mesh.point(vh2), mesh.point(vh3), mesh.point(vh4)

            S123 = 0.5 * np.linalg.norm(np.cross(p2 - p1, p2 - p3))
            S134 = 0.5 * np.linalg.norm(np.cross(p4 - p1, p4 - p3))
            l13 = np.linalg.norm(p1 - p3)

            coef1 = (S123 * np.dot(p4 - p3, p3 - p1) + S134 * np.dot(p1 - p3, p3 - p2)) / (l13 * l13 * (S123 + S134))
            coef2 = S134 / (S123 + S134)
            coef3 = (S123 * np.dot(p3 - p1, p1 - p4) + S134 * np.dot(p2 - p1, p1 - p3)) / (l13 * l13 * (S123 + S134))
            coef4 = S123 / (S123 + S134)

            D_row.append(cnt)
            D_col.append(vh1.idx())
            D_data.append(coef1)

            D_row.append(cnt)
            D_col.append(vh2.idx())
            D_data.append(coef2)

            D_row.append(cnt)
            D_col.append(vh3.idx())
            D_data.append(coef3)

            D_row.append(cnt)
            D_col.append(vh4.idx())
            D_data.append(coef4)

            R_data.append(1)
            R_data.append(-1)
            R_data.append(1)
            R_data.append(-1)

        D = coo_matrix((D_data, (D_row, D_col)), shape=(N_edge, N_vert))
        R = coo_matrix((R_data, (D_row, D_col)), shape=(N_edge, N_vert))
        p = np.array([mesh.point(vertex) for vertex in mesh.vertices()])
        p_star = copy.deepcopy(p)
        return D, R, p, p_star

    def optimize(self) -> TriMesh:
        mesh = self.mesh
        N_edge = mesh.n_edges()
        N_vert = mesh.n_vertices()

        if self.type == 'edge':
            N = N_edge
        else: # type=='vert
            N = N_vert

        D, R, p, p_star = self.build_operator()
        start_time = time.time()
        # --------------build linear system--------------
        cnt = 1
        while self.beta < self.beta_max:
            # stage 1:
            delta = D @ p

            for i in range(N):
                if delta[i][0]**2 + delta[i][1]**2 + delta[i][2]**2 < self.lamba / self.beta:
                    delta[i][0] = delta[i][1] = delta[i][2] = 0.0

            # stage 2:
            A = sp.eye(N_vert) + self.alpha * (R.transpose() @ R) + self.beta * (D.transpose() @ D)
            b = p_star + self.beta * (D.transpose() @ delta)

            if self.t_solver == 'ch':
                factor = cholesky(A)
            # p, info = sp.linalg.cg(A, b, x0=p)
            for i in range(3):
                if self.t_solver == 'cg':
                    p[:,i], info = sp.linalg.cg(A, b[:,i], x0=p[:,i])
                if self.t_solver == 'ch':
                    # scipy cholesky
                    # y = scipy.linalg.solve_triangular(L, b[:,i], lower=True, check_finite=False)
                    # p[:,i] = scipy.linalg.solve_triangular(L.T, y, lower=True, check_finite=False)
                    p[:,i] = factor.solve_A(b[:,i])
            if self.log:
                print('iter {}: beta={}, p[0,0]={}'.format(cnt, self.beta, p[0][0]))
            cnt += 1
            self.beta *= self.kappa
            self.alpha *= 0.5
        self.denoised_vert = p
        final_time = time.time()
        self.optimize_time = final_time - start_time
        self.iteration = cnt
        # --------------return mesh--------------
        for i, vh in enumerate(mesh.vertices()):
            mesh.set_point(vh, p[i])
        self.denoised_mesh = mesh
        return mesh

if __name__=='__main__':
    parser = argparse.ArgumentParser(description="implementation of image smoothing via L0 gradient minimization")

    #parser.add_argument('--input', default='./input/fandisk.obj',
    parser.add_argument('--input', default='./input/SharpSphere.obj',
        help="input mesh file")

    parser.add_argument('--output', default='./res.obj',
        help="output mesh file")

    parser.add_argument('-k', type=float, default=1.414,
        metavar='kappa', help='updating weight (default 1.414), refer mu in mesh2013')

    parser.add_argument('-l', type=float, default=1./16,
        metavar='lambda', help='smoothing weight (None is automatic)')

    parser.add_argument('--beta', type=float, default=1.0e-3,
        help='beta init value (default 1.0e-3)')

    parser.add_argument('--beta_max', type=float, default=1e3,
        help='updating threshold (default 1e3)')

    parser.add_argument('-v', '--verbose', action='store_true',
        help='enable verbose logging for each iteration')

    parser.add_argument('-m', default='edge',
        metavar='method', help="support three method: vert, edge, area")

    parser.add_argument('-n', type=float, default=0.3,
        metavar='noise', help="ratio of average edge length")

    parser.add_argument('-r', '--regulation', action='store_true',
        help='use regulation')

    parser.add_argument('-a', type=float, default=0,
        metavar='alpha', help="parameter of regulation")

    args = parser.parse_args()

    #s = LinearSystem_Solver(args)
    if args.m == 'vert':
        s = VertexSolver(args)
    elif args.m == 'edge':
        s = CotEdgeSolver(args)
    elif args.m == 'area':
        if args.regulation:
            s = AreaEdgeSolver_R(args)
        else:
            s = AreaEdgeSolver(args)
    elif args.m == 'vfft':
        s = VertexSolver_FFT(args)

    noise_mesh = add_noise(s.input_mesh, l=float(args.n))
    s.save(noise_mesh, './noise.obj')
    s.prepare_mesh(mesh=noise_mesh)
    s.solve()
    s.save(s.denoised_mesh)

    print('done!')


# if __name__ == '__main__':
#     str = input("Enter the string: ")

#     print("Characters of the string are: ")
#     set1 = set(str)
#     for i in set1:
#         print(f'{i}', end=' ')

#     if(str.isalpha()):
#         if(not str.isupper()):
#             print(f"\n\nUppercase of {str} = {str.upper()}")
#         if(not str.islower()):
#             print(f"Lowercase of {str} = {str.lower()}")
#     else:
#         print("\nCannot perform string operations.")

#     str1 = input("\nEnter the string for concatination: ")
#     print(f"{str} + {str1} = {str+' '+str1}")

# n = 5

# for row in range(1, n+1):
#     for star in range(1, row+1):
#         print("*", end='')
#     print()

n = 13

for i in range(2, n):
    if(n % i == 0):
        print("Not a prime")
        break
else:
    print("Prime")


from django.core.management import BaseCommand, CommandError
from chains.models import Story, Chain
import json
import glob
import os

PROTEST_PATH = '/Users/sneha/Documents/dev/narrate.ai/narratives/output/protest_events/'
protest_files = glob.glob(os.path.join(PROTEST_PATH, '*.json'))
CHAIN_PATH = '/Users/sneha/Documents/dev/narrate.ai/narratives/output/chains_dj_wapo/'
files = glob.glob(os.path.join(CHAIN_PATH, '*.json'))

class Command(BaseCommand):
    help = 'Run this command to populate db'

    def handle(self, *args, **options):
        # for p_file in protest_files:
        #     with open(p_file, 'r') as f:
        #         for line in f:
        #             data = json.loads(line)
        #             s = Story(id=data['StoryID'], text=data['RawText'],
        #                       publisher=data['Publisher'], event_code=int(data['event_details']['CAMEO Code']),
        #                       country=data['event_details']['Country'])
        #             s.save()
        #             print(data['StoryID'])

        for c_file in files:
            story_id = int(c_file.split('/')[-1].split('.')[0])
            with open(c_file, 'r') as f:
                chains = json.load(f)

            c = Chain(story_id=Story.objects.get(id=story_id), chains=str(chains))
            c.save()

            print(story_id)


# 내리막 길
from sys import stdin
from collections import deque

m, n = map(int, stdin.readline().rstrip().split())
graph = [list(map(int, stdin.readline().rstrip().split())) for _ in range(m)]

result = [[-1] * n for _ in range(m)]
dx = [0, 0, -1, 1]
dy = [-1, 1, 0, 0]

def calc(result, table, m, n, x, y):
    if x == m - 1 and y == n - 1:
        result[x][y] = 1
        return
    if result[x][y] != -1:
        return
    result[x][y] = 0
    for i in range(4):
        nx = x + dx[i]
        ny = y + dy[i]
        if 0 <= nx < m and 0 <= ny < n:
            if graph[x][y] > graph[nx][ny]:
                calc(result, graph, m, n, nx, ny)
                result[x][y] += result[nx][ny]

calc(result, graph, m, n, 0, 0)
print(result[0][0])


def cria_lista_multiplos(numero):
    if (numero % 1) != 0 or numero <= 0:
        return 'Número inválido!'
    else:
        tuplo = ()
        countador = 0
        indice = 0
        while countador < 10:
            if (indice % numero) == 0:
                countador += 1
                tuplo += (indice,)
            indice += 1
        return tuplo

num_intr = eval(input('Insira um número inteiro positivo: '))
print(cria_lista_multiplos(num_intr))


from lxml import html
import requests
import pandas as pd

page_tennis = requests.get('http://www.espn.com/tennis/players')
tree_tennis = html.fromstring(page_tennis.content)

p = tree_tennis.cssselect(".tablehead a")

players = []
for player in p:
    player = player.text_content()
    players.append(player)

c = tree_tennis.cssselect("td + td")

countries = []
for country in  c:
        country = country.text_content()
        if len(country) == 0:
            countries.append("Null")
        else:
            countries.append(country)

for alll in countries:
    if "COUNTRY" in countries: countries.remove("COUNTRY")

df_tennis = pd.DataFrame(list(zip(players, countries)), columns = ['Name', 'Country'])

page_ranking = requests.get("https://www.espn.com/tennis/rankings")
tree_ranking = html.fromstring(page_ranking.content)

r = tree_ranking.cssselect(".Table__TD , .mr7")

rankings = []
for i in r:
    i = i.text_content()
    rankings.append(i)

#table_tree = tree_tennis.xpath('//td')
#z=[]
#for t in table_tree:
#    t = t.text_content()
#    z.append(t)
#
#col=[table_tree[1].text_content(), table_tree[2].text_content()]
#
#players=[]
#for t in table_tree:
#    player = t.text_content()#[3:len(table_tree):2]
#    players.append(player)
#
#players = players[3:len(table_tree):2]
#for player in players:
#    player =
#
#
#i=0
#for t in table_tree[1]:
#    i+= 1
#    name = t.text_content()
#    #print(name)
#    print('%d:"%s"'%(i,name))
#    col.append((name,[]))
#
#for j in range(1,len(table_tree)):
#    T=table_tree[j]
#    if len(T)!=2:
#        break
#    i=0
#    for t in T.iterchildren():
#        data=t.text_content().strip()
#        col[i][1].append(data)
#        i+=1
#
#Dict_tennis = {title:column for (title,column) in col}


import glob
import random
import os
import sys
import numpy as np
from PIL import Image
import torch
import torch.nn.functional as F

from pycocotools.coco import COCO
from utils.augmentations import horisontal_flip
from torch.utils.data import Dataset
import torchvision.transforms as transforms

def pad_to_square(img, pad_value):
    c, h, w = img.shape
    dim_diff = np.abs(h - w)
    # (upper / left) padding and (lower / right) padding
    pad1, pad2 = dim_diff // 2, dim_diff - dim_diff // 2
    # Determine padding
    pad = (0, 0, pad1, pad2) if h <= w else (pad1, pad2, 0, 0)
    # Add padding
    img = F.pad(img, pad, "constant", value=pad_value)

    return img, pad

def resize(image, size):
    image = F.interpolate(image.unsqueeze(0), size=size, mode="nearest").squeeze(0)
    return image

def random_resize(images, min_size=288, max_size=448):
    new_size = random.sample(list(range(min_size, max_size + 1, 32)), 1)[0]
    images = F.interpolate(images, size=new_size, mode="nearest")
    return images

class ImageFolder(Dataset):

    def __init__(self, ann_file='annotations/instances_train2017.json', data_root='data/coco', img_size=416, mode='train'):
        self.ann_file = ann_file
        self.data_root = data_root

        if self.data_root is not None:
            if not os.path.isabs(self.ann_file):
                self.ann_file = os.path.join(self.data_root, self.ann_file)

        self.data_infos = self.load_annotations(self.ann_file)
        if mode == 'train':
            self.files = [self.data_root + '/train2017/' + self.data_infos[i]['filename'] for i in range(len(self.data_infos))]
        elif mode == 'val':
            self.files = [self.data_root + '/val2017/' + self.data_infos[i]['filename'] for i in range(len(self.data_infos))]
        else:
            self.files = [self.data_root + '/test2017/' + self.data_infos[i]['filename'] for i in range(len(self.data_infos))]
        self.img_size = img_size

    def __getitem__(self, index):
        img_path = self.files[index % len(self.files)]
        # Extract image as PyTorch tensor
        img = transforms.ToTensor()(Image.open(img_path))
        # Pad to square resolution
        img, _ = pad_to_square(img, 0)
        # Resize
        img = resize(img, self.img_size)

        return img_path, img

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

    def load_annotations(self, ann_file):
        self.coco = COCO(ann_file)
        self.cat_ids = self.coco.get_cat_ids(cat_names=self.CLASSES)
        self.cat2label = {cat_id: i for i, cat_id in enumerate(self.cat_ids)}
        self.img_ids = self.coco.get_img_ids()
        data_infos = []
        for i in self.img_ids:
            info = self.coco.load_imgs([i])[0]
            info['filename'] = info['file_name']
            data_infos.append(info)
        return data_infos

class CocoDataset(Dataset):

    CLASSES = ('person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
               'train', 'truck', 'boat', 'traffic light', 'fire hydrant',
               'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog',
               'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe',
               'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
               'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat',
               'baseball glove', 'skateboard', 'surfboard', 'tennis racket',
               'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl',
               'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot',
               'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
               'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop',
               'mouse', 'remote', 'keyboard', 'cell phone', 'microwave',
               'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock',
               'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush')

    def __init__(self,
                 ann_file='annotations/instances_train2017.json',
                 data_root='data/coco',
                 img_size=416,
                 mode='train',
                 augment=True,
                 multiscale=True,
                 normalized_labels=False):
        self.ann_file = ann_file
        self.data_root = data_root

        if self.data_root is not None:
            if not os.path.isabs(self.ann_file):
                self.ann_file = os.path.join(self.data_root, self.ann_file)

        self.data_infos = self.load_annotations(self.ann_file)
        if mode == 'train':
            self.files = [self.data_root + '/train2017/' + self.data_infos[i]['filename'] for i in range(len(self.data_infos))]
        elif mode == 'val':
            self.files = [self.data_root + '/val2017/' + self.data_infos[i]['filename'] for i in range(len(self.data_infos))]
        else:
            self.files = [self.data_root + '/test2017/' + self.data_infos[i]['filename'] for i in range(len(self.data_infos))]

        self.img_size = img_size
        self.max_objects = 100
        self.augment = augment
        self.multiscale = multiscale
        self.normalized_labels = normalized_labels
        self.min_size = self.img_size - 3 * 32
        self.max_size = self.img_size + 3 * 32
        self.batch_count = 0

    def __getitem__(self, index):

        # ---------
        #  Image
        # ---------

        img_path = self.files[index]

        # Extract image as PyTorch tensor
        img = transforms.ToTensor()(Image.open(img_path).convert('RGB'))

        # Handle images with less than three channels
        if len(img.shape) != 3:
            img = img.unsqueeze(0)
            img = img.expand((3, img.shape[1:]))

        _, h, w = img.shape
        h_factor, w_factor = (h, w) if self.normalized_labels else (1, 1)
        # Pad to square resolution
        img, pad = pad_to_square(img, 0)
        _, padded_h, padded_w = img.shape

        # ---------
        #  Label
        # ---------

        targets = None
        coco_id_map={1: 0, 2: 1, 3: 2, 4: 3, 5: 4, 6: 5, 7: 6, 8: 7, 9: 8, 10: 9, 11: 10,
                     13: 11, 14: 12, 15: 13, 16: 14, 17: 15, 18: 16, 19: 17, 20: 18, 21: 19, 22: 20,
                     23: 21, 24: 22, 25: 23, 27: 24, 28: 25, 31: 26, 32: 27, 33: 28, 34: 29, 35: 30,
                     36: 31, 37: 32, 38: 33, 39: 34, 40: 35, 41: 36, 42: 37, 43: 38, 44: 39, 46: 40,
                     47: 41, 48: 42, 49: 43, 50: 44, 51: 45, 52: 46, 53: 47, 54: 48, 55: 49, 56: 50,
                     57: 51, 58: 52, 59: 53, 60: 54, 61: 55, 62: 56, 63: 57, 64: 58, 65: 59, 67: 60,
                     70: 61, 72: 62, 73: 63, 74: 64, 75: 65, 76: 66, 77: 67, 78: 68, 79: 69, 80: 70,
                     81: 71, 82: 72, 84: 73, 85: 74, 86: 75, 87: 76, 88: 77, 89: 78, 90: 79}
        img_id = self.data_infos[index]['id']
        ann_ids = self.coco.get_ann_ids(img_ids=[img_id])
        ann_info = self.coco.load_anns(ann_ids)
        cat_ids = [coco_id_map[ann['category_id']] for ann in ann_info]
        bboxes = [ann['bbox'] for ann in ann_info]

        cat_ids = torch.from_numpy(np.array(cat_ids))
        bboxes = torch.from_numpy(np.array(bboxes))

        if len(bboxes.size()) == 2:
            # Extract coordinates for unpadded + unscaled image
            x1 = w_factor * (bboxes[:, 0] - bboxes[:, 2] / 2)
            y1 = h_factor * (bboxes[:, 1] - bboxes[:, 3] / 2)
            x2 = w_factor * (bboxes[:, 0] + bboxes[:, 2] / 2)
            y2 = h_factor * (bboxes[:, 1] + bboxes[:, 3] / 2)
            # Adjust for added padding
            x1 += pad[0]
            y1 += pad[2]
            x2 += pad[1]
            y2 += pad[3]
            # Returns (x, y, w, h)
            bboxes[:, 0] = ((x1 + x2) / 2) / padded_w
            bboxes[:, 1] = ((y1 + y2) / 2) / padded_h
            bboxes[:, 2] *= w_factor / padded_w
            bboxes[:, 3] *= h_factor / padded_h

            targets = torch.zeros((len(bboxes), 6))
            targets[:, 1] = cat_ids
            targets[:, 2:] = bboxes

            # Apply augmentations
            if self.augment:
                if np.random.random() < 0.5:
                    img, targets = horisontal_flip(img, targets)

        return img_path, img, targets

    def collate_fn(self, batch):
        paths, imgs, targets = list(zip(*batch))
        # Remove empty placeholder targets
        targets = [boxes for boxes in targets if boxes is not None]
        # Add sample index to targets
        for i, boxes in enumerate(targets):
            boxes[:, 0] = i
        targets = torch.cat(targets, 0)
        # Selects new image size every tenth batch
        if self.multiscale and self.batch_count % 10 == 0:
            self.img_size = random.choice(range(self.min_size, self.max_size + 1, 32))
        # Resize images to input shape
        imgs = torch.stack([resize(img, self.img_size) for img in imgs])
        self.batch_count += 1
        return paths, imgs, targets

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

    def load_annotations(self, ann_file):
        self.coco = COCO(ann_file)
        self.cat_ids = self.coco.get_cat_ids(cat_names=self.CLASSES)
        self.cat2label = {cat_id: i for i, cat_id in enumerate(self.cat_ids)}
        self.img_ids = self.coco.get_img_ids()
        data_infos = []
        for i in self.img_ids:
            info = self.coco.load_imgs([i])[0]
            info['filename'] = info['file_name']
            data_infos.append(info)
        return data_infos


#!/usr/bin/env python3

import sys
import random
from PyQt5 import QtWidgets, QtCore, QtGui

class TegenBoard(QtWidgets.QWidget):
    """
    Tegen board widget. Currently has the Tegen grid.
    Inherits from QWidget
    Nothing here is for external use.

    Overwritten or new methods:
    void init_ui(): defines and draws the widget
    void paintEvent(QPaintEvent e): redraws the widget on each paint event
    void resizeEvent(QResizeEvent e): redefines the grid when widget is
                                      resized and queues it to be redrawn
    void keyPressEvent(QKeyEvent e): When a key is pressed, sets a single tile
                                     to a random color. The tile that is
                                     recolored is determined by taking the
                                     modulo of the keycode and the number of
                                     tiles.

    Instance variables:
    dict tiles[] = {"rect": QRect, "color": QColor, "image": QImage}:
        List of arrays in the described form. Keys are as follows:
        QRect rect: defines the object to paint. Initially set by init_ui, and
                    updated by resizeEvent()
        QColor color: the fill color. Initially set by init_ui() to
                      QColor(255, 255, 255), and updated to a random color by
                      keyPressEvent.
        QImage image: Initially set to None. Should this be set to an actual
                      value, it will be painted on the tile.
    """
    def __init__(self):
        super().__init__()
        self.init_ui()

    def init_ui(self):
        self.setWindowTitle("Tegen")

        self.tiles = []

        for i in range(7):
            for j in range(7):
                if abs(i - 3) + abs(j - 3) >= 5:
                    continue

                self.tiles.append({"rect": QtCore.QRect((self.width() / 7) * i,
                    (self.height() / 7) * j, self.width() / 7 - 1,
                    self.height() / 7 - 1),
                    "color": QtGui.QColor(255,255,255),
                    "image": None
                })

        self.show()

    def paintEvent(self, e):
        painter = QtGui.QPainter(self)
        painter.setPen(QtGui.QColor(0,0,0))
        for i in range(len(self.tiles)):
            painter.drawRect(self.tiles[i]["rect"])
            painter.fillRect(self.tiles[i]["rect"], self.tiles[i]["color"])
            if self.tiles[i]["image"]:
                painter.drawImage(self.tiles[i]["rect"], self.tiles[i]["image"])

    def resizeEvent(self, e):
        idx = 0

        for i in range(7):
            for j in range(7):
                if abs(i - 3) + abs(j - 3) >= 5:
                    continue

                self.tiles[idx]["rect"].setRect((self.width() / 7) * i,
                    (self.height() / 7) * j, self.width() / 7 - 1,
                    self.height() / 7 - 1)

                idx += 1

        self.update()

    def keyPressEvent(self, e):
        self.tiles[e.key() % len(self.tiles)]["color"] = QtGui.QColor(
            *[random.randint(0,255) for i in range(3)])

        self.repaint(self.tiles[e.key() % len(self.tiles)]["rect"])

def main():
    app = QtWidgets.QApplication(sys.argv)
    board = TegenBoard()
    sys.exit(app.exec_())

if __name__ == '__main__':
    main()


import Playfield
import Pacman
import Candy

class Game:
    """
    La classe Game initialise le jeu

    Attributes
    ----------
    gamestate : bool
    indique l'état du jeu (en cours ou non), false, le jeu n'est pas en cours, true le jeu est en cours

    Methods
    -------
    Newgame : BOOL
        lance le jeu et instancie les éléments nécessaires
    endofgame : BOOL
        regarde si nicepacman ou badpacman est mort et termine le jeu si c'est le cas
    intercept_candy :
        regarde si un bonbon est rencontré lors du déplacement d'un pacman
    intercept_wall :
        regarde si un mur est rencontré lors du déplacement d'un pacman

    """
    def __init__(self, gamestate):

        """
        Initialise le jeu

        Attributes
        ----------
        gameState = indique l'état du jeu (en cours ou non), false, le jeu n'est pas en cours, true le jeu est en cours

        """
        self.nicepacman1 = Pacman.NicePacman([7,8],0)
        self.badpacman1 = Pacman.BadPacman([8,8],100)
        self.terrain = Playfield.Playfield()
        self.gamestate = gamestate

    def new_game(self):
        """
        Lance le jeu et et créer les éléments nécessaires

        """
        start=str(input("Tapez GO pour commencer "))
        if start=="Go":
            print("Ok, let's begin!")
            self.terrain.create_matrice()
            self.terrain.create_candies()
            self.nicepacman1.show_parameters
            self.badpacman1.show_parameters
            #return gamestate == True

    def tour(self):
        """
        Définit comment un tour du jeu se passe

        Returns
        -------
        Retourne les paramètres de chaque pacman

        """
        while self.gamestate==True:
            print ("déplacement Nicepacman")
            print ("déplacement Badpacman")
            self.nicepacman1.move()
            self.badpacman1.move()

            if self.terrain.is_wall(self.nicepacman1.position) == True:
                print("Oups Nice, c'est un mur, tu dois de déplacer à nouveau !")
                self.nicepacman1.move()
            elif self.terrain.is_wall(self.badpacman1.position) == True:
                print("Oups Bad, c'est un mur, tu dois de déplacer à nouveau !")
                self.badpacman1.move()
            else:
                self.nicepacman1.eat()
                self.badpacman2.eat()

            for x in (self.terrain.create_candies):
                if self.nicepacman1.position == self.terrain.create_candies[x].position:
                    self.nicepacman1.eat(self.terrain.create_candies[x])
                elif self.badpacman1.position == self.terrain.create_candies[x].position:
                    self.badpacman1.eat(self.terrian.create_candies[x])

    def end_of_game(self):
            """
            Définit comment le jeu se termine

            Returns
            -------
            Retourne l'état du jeu, True si rien ne se passe, False si un des pacman est mort

            """
            if (self.badpacman1.is_dead() == True) or (self.nicepacman1.is_dead() == True):
                self.gamestate = False
                return self.gamestate
            else:
                self.gamestate = True
                return self.gamestate


#!/usr/bin/env python3
from argparse import ArgumentParser

from gooey import Gooey
from pyswip import Prolog

swipl = Prolog()
swipl.consult("trains.pl")

STATIONS = [ans["S"] for ans in swipl.query("station(S)")]

def route_length(From, To):
    return next(swipl.query(f"route_length({From}, {To}, L)"))["L"]

@Gooey
def main():
    parser = ArgumentParser(description="Find the length of a route between two stations")
    parser.add_argument("-o", "--origin", choices=STATIONS, required=True)
    parser.add_argument("-d", "--destination", choices=STATIONS, required=True)
    args = parser.parse_args()
    print(route_length(args.origin, args.destination))

if __name__ == "__main__":
    main()