import logging
from typing import Optional, List, Dict, Any, Type
import re

from enums.query.query_output_types import QueryOutputTypes
from enums.query.props.prop_enum import PropEnum
from enums.cloud_domains import CloudDomains
from enums.query.query_presets import (
    QueryPresetsString,
    QueryPresetsGeneric,
    QueryPresetsDateTime,
)

from structs.query.query_output_details import QueryOutputDetails
from structs.query.query_preset_details import QueryPresetDetails

from openstack_query.api.query_api import QueryAPI
from custom_types.openstack_query.aliases import QueryReturn

from exceptions.enum_mapping_error import EnumMappingError

logger = logging.getLogger(__name__)

# pylint:disable=too-few-public-methods

# these search methods map to stackstorm actions which just pass
# user-inputted strings straight here - hence there's a lot of parameters and local variables
# pylint:disable=too-many-arguments
# pylint:disable=too-many-locals
class ManagerWrapper:
    """
    This class is the base class for all managers.
    Managers hold a query object and several methods which build and run specific queries
    """

    def __init__(
        self, query: QueryAPI, cloud_account: CloudDomains, prop_cls: Type[PropEnum]
    ):
        self._query = query
        self._cloud_account = cloud_account
        self._prop_cls = prop_cls

    def _build_and_run_query(
        self,
        output_details: QueryOutputDetails,
        preset_details: Optional[QueryPresetDetails] = None,
        runner_params: Optional[Dict[str, Any]] = None,
    ) -> QueryReturn:
        """
        method to build the query, execute it, and return the results
        :param output_details: A dataclass containing config on how to output results of query
        :param preset_details: An optional dataclass containing query preset config information
        :param runner_params: An optional set of extra params to pass when calling run()
        """
        logging.info("Running Query")
        if preset_details:
            preset_args = "\n\t".join(
                [f"{key}: '{val}'" for key, val in preset_details.args.items()]
            )
            log_preset_out = f"{preset_details.prop.name} {preset_details.preset.name} with args {preset_args}"
            logging.info("Query Details: %s", log_preset_out)
        else:
            logging.info("Query Details: Getting All")

        if output_details.output_type == QueryOutputTypes.TO_OBJECT_LIST:
            # properties to select will be ignored - hence don't log it
            logging.info(
                "Output Details:\n\t output_type: %s", output_details.output_type.name
            )
        else:
            props_to_select = ",".join(
                [prop.name for prop in output_details.properties_to_select]
            )
            logging.info(
                "Outputting as:\n\t output_type: %s\n\t with properties: %s",
                output_details.output_type.name,
                props_to_select,
            )

        if not runner_params:
            runner_params = {}

        self._populate_output_params(output_details=output_details)

        if preset_details:
            self._query.where(
                preset=preset_details.preset,
                prop=preset_details.prop,
                **preset_details.args,
            )

        self._query.run(self._cloud_account, **runner_params)

        return self._get_query_output(
            output_details.output_type,
        )

    def _populate_output_params(self, output_details: QueryOutputDetails):
        """
        method that sets output parameters like group by, sort by and select parameters
        :param output_details: A dataclass containing output parameters to use
        """

        # set select props
        self._query.select(*output_details.properties_to_select)

        # set sort by
        if output_details.sort_by:
            self._query.sort_by(*output_details.sort_by)

        # set group by
        if output_details.group_by:
            self._query.group_by(
                group_by=output_details.group_by,
                group_ranges=output_details.group_ranges,
                include_ungrouped_results=output_details.include_ungrouped_results,
            )

    def _get_query_output(
        self,
        output_type: QueryOutputTypes,
    ) -> QueryReturn:
        """
        method that returns the output of query
        :param output_type: An Enum representing how to output query results
        """
        output_func = {
            QueryOutputTypes.TO_STR: self._query.to_string,
            QueryOutputTypes.TO_HTML: self._query.to_html,
            QueryOutputTypes.TO_LIST: self._query.to_list,
            QueryOutputTypes.TO_OBJECT_LIST: lambda: self._query.to_list(
                as_objects=True
            ),
        }.get(output_type, None)
        if not output_func:
            logging.error(
                "Error: No function mapping found for output type %s "
                "- if you are here as a developer, you must add a mapping in this method"
                "for the enum to a public function in QueryMethods",
                output_type.name,
            )
            raise EnumMappingError(
                f"could not find a mapping for output type {output_type.name}"
            )
        return output_func()

    def search_all(
        self,
        properties_to_select: Optional[List[str]] = None,
        output_type: Optional[str] = None,
        group_by: Optional[str] = None,
        sort_by: Optional[List[str]] = None,
        **kwargs,
    ) -> QueryReturn:
        """
        method that returns a list of all resources
        :param properties_to_select: list of strings representing which properties to select
        :param output_type: string representing how to output the query
        :param group_by: an optional string representing a property to group results by
        :param sort_by: an optional set of tuples representing way which properties to sort results by
        :param kwargs: A set of optional meta params to pass to the query

        """
        logging.info("Running 'search all' query - will output all values")
        output_details = QueryOutputDetails.from_kwargs(
            prop_cls=self._prop_cls,
            properties_to_select=properties_to_select,
            output_type=output_type,
            group_by=group_by,
            sort_by=sort_by,
        )

        return self._build_and_run_query(
            preset_details=None, output_details=output_details, runner_params=kwargs
        )

    def search_by_property(
        self,
        search_mode: str,
        property_to_search_by: str,
        values: List[str],
        properties_to_select: Optional[List[str]] = None,
        output_type: Optional[str] = None,
        group_by: Optional[str] = None,
        sort_by: Optional[List[str]] = None,
        **kwargs,
    ) -> QueryReturn:
        """
        method that builds and runs a query to find generic resource with a selected property
        matching, or not matching given value(s)
        :param search_mode: A string representing a preset Enum ANY_IN or NOT_ANY_IN which dictates what query
        to perform
        if False - use the preset NOT_ANY_IN/NOT_EQUAL_TO
        :param property_to_search_by: A string representing a datetime property Enum that the preset will be used on
        :param values: A list of string values to compare server property against
        :param properties_to_select: list of strings representing which properties to select
        :param output_type: string representing how to output the query
        :param group_by: an optional string representing a property to group results by
        :param sort_by: an optional set of tuples representing way which properties to sort results by
        :param kwargs: A set of optional meta params to pass to the query
        """
        # convert user-given args into enums
        logging.info("Running search by property query")
        logging.debug("converting user-defined property string into enum")
        preset = (
            QueryPresetsGeneric.ANY_IN
            if search_mode
            else QueryPresetsGeneric.NOT_ANY_IN
        )
        prop = self._prop_cls.from_string(property_to_search_by)
        args = {"values": values}

        logging.info(
            "This query will find all resources for which %s matches any of [%s]",
            prop.name,
            args["values"],
        )

        # If values contains only one value - use EQUAL_TO/NOT_EQUAL_TO as the preset instead to speed up query
        if len(values) == 1:
            logging.info(
                "query only contains one value - EQUAL_TO preset will be used to speed up query"
            )
            equal_to_preset = {
                QueryPresetsGeneric.ANY_IN: QueryPresetsGeneric.EQUAL_TO,
                QueryPresetsGeneric.NOT_ANY_IN: QueryPresetsGeneric.NOT_EQUAL_TO,
            }.get(preset, None)
            if equal_to_preset:
                preset = equal_to_preset
                args = {"value": values[0]}
        else:
            logging.info("query contains multiple values - ANY_IN preset will be used")

        output_details = QueryOutputDetails.from_kwargs(
            prop_cls=self._prop_cls,
            properties_to_select=properties_to_select,
            output_type=output_type,
            group_by=group_by,
            sort_by=sort_by,
        )

        return self._build_and_run_query(
            preset_details=QueryPresetDetails(
                preset=preset,
                prop=prop,
                args=args,
            ),
            output_details=output_details,
            runner_params=kwargs,
        )

    def search_by_regex(
        self,
        property_to_search_by: str,
        pattern: str,
        properties_to_select: Optional[List[str]] = None,
        output_type: Optional[str] = None,
        group_by: Optional[str] = None,
        sort_by: Optional[List[str]] = None,
        **kwargs,
    ):
        """
        method that builds and runs a query to find generic resource with a property matching regex.
        :param property_to_search_by: A string representing a string property Enum that the preset will be used on
        :param pattern: A string representing a regex pattern
        :param properties_to_select: list of strings representing which properties to select
        :param output_type: string representing how to output the query
        :param group_by: an optional string representing a property to group results by
        :param sort_by: an optional set of tuples representing way which properties to sort results by
        :param kwargs: A set of optional meta params to pass to the query
        """
        logging.info("Running search by property query")

        logging.debug("converting user-defined property string into enum")
        prop = self._prop_cls.from_string(property_to_search_by)

        logging.debug("checking that user-defined regex pattern is valid")
        re.compile(pattern)

        args = {"regex_string": pattern}

        logging.info(
            "This query will find all resources for which %s matches a regex pattern: '%s'",
            prop.name,
            pattern,
        )

        output_details = QueryOutputDetails.from_kwargs(
            prop_cls=self._prop_cls,
            properties_to_select=properties_to_select,
            output_type=output_type,
            group_by=group_by,
            sort_by=sort_by,
        )

        return self._build_and_run_query(
            preset_details=QueryPresetDetails(
                preset=QueryPresetsString.MATCHES_REGEX,
                prop=prop,
                args=args,
            ),
            output_details=output_details,
            runner_params=kwargs,
        )

    def search_by_datetime(
        self,
        search_mode: str,
        property_to_search_by: str,
        days: int = 0,
        hours: int = 0,
        minutes: int = 0,
        seconds: int = 0,
        properties_to_select: Optional[List[str]] = None,
        output_type: Optional[str] = None,
        group_by: Optional[str] = None,
        sort_by: Optional[List[str]] = None,
        **kwargs,
    ) -> QueryReturn:
        """
        method that builds and runs a datetime-related query on Generic Resource, and then returns results.
        Uses UTC timezone
        :param search_mode: A string representing what type of datetime query will be run - set as a Datetime Preset
        that the query will use
        :param property_to_search_by: A string representing a datetime property Enum that the preset will be used on
        - must be datetime compatible
        :param days: (Optional) Number of relative days in the past from now to use as threshold
        :param hours: (Optional) Number of relative hours in the past from now to use as threshold
        :param minutes: (Optional) Number of relative minutes in the past from now to use as threshold
        :param seconds: (Optional) Number of relative seconds in the past from now to use as threshold
        :param properties_to_select: list of strings representing which properties to select
        :param output_type: string representing how to output the query
        :param group_by: an optional string representing a property to group results by
        :param sort_by: an optional set of tuples representing way which properties to sort results by
        :param kwargs: A set of optional meta params to pass to the query
        """
        logging.info("Running search by property query")

        logging.debug("converting user-defined property string into enum")
        prop = self._prop_cls.from_string(property_to_search_by)
        preset = QueryPresetsDateTime.from_string(search_mode)
        args = {
            "days": days,
            "hours": hours,
            "minutes": minutes,
            "seconds": seconds,
        }
        logging.info(
            "This query will run a relative datetime query %s %s with args %s",
            prop.name,
            preset.name,
            ",".join(f"{key}: '{val}'" for key, val in args.items()),
        )

        preset_details = QueryPresetDetails(
            preset=preset,
            prop=prop,
            args=args,
        )

        output_details = QueryOutputDetails.from_kwargs(
            prop_cls=self._prop_cls,
            properties_to_select=properties_to_select,
            output_type=output_type,
            group_by=group_by,
            sort_by=sort_by,
        )

        return self._build_and_run_query(
            preset_details=preset_details,
            output_details=output_details,
            runner_params=kwargs,
        )


from astropy.table import Table
from desitarget.targetmask import desi_mask, bgs_mask, mws_mask
import numpy as np
mtl = Table.read('mtl.fits')
zcat = Table.read("zcat.fits")
print("Efficiencies N_in N_out N_in/N_out")
for objtype in ['BGS_ANY', 'ELG', 'LRG', 'QSO', 'STD_FSTAR', 'STD_BRIGHT']:
    ii = (mtl['DESI_TARGET'] & desi_mask.mask(objtype)) != 0
    jj = np.in1d(mtl['TARGETID'][ii], zcat['TARGETID'])
    n_in = np.count_nonzero(ii)
    n_out = np.count_nonzero(jj)
    print('{:10} {} {} {:.2f}'.format(objtype, n_in, n_out, n_out/n_in))


import random
num1=random.randint(0,100)
num2=random.randint(0,100)
print(num1,num2)
t=0
if num1<num2:
   t=num1
   num1=num2
   num2=t
print("what is num1-num2:")
answer=eval(input("enter your answer"))
if answer==num1-num2:
    print("correct")
else:
    print("wrong")
count=eval(input("enter count"))
print(bool(count%10==0))


from django.contrib import admin
from django.urls import path, include
from django.conf import settings
from django.conf.urls.static import static
from django.conf.urls.i18n import i18n_patterns
from . import views

urlpatterns = [
    path('', views.select_lang_redirect),
    path('captcha/', include('captcha.urls')),
]

urlpatterns += i18n_patterns(
    path('', include('MW_Website.urls', namespace='website')),
    path('', include('MW_Auth.urls', namespace='auth')),
    path('', include('MW_Setting.urls', namespace='setting')),
    path('change/language/', views.activate_language, name='activate_lang'),
    path('admin/', admin.site.urls),
)

if settings.DEBUG:
    urlpatterns = urlpatterns + static(settings.STATIC_URL, document_root=settings.STATIC_ROOT)
    urlpatterns = urlpatterns + static(settings.MEDIA_URL, document_root=settings.MEDIA_ROOT)


# Definition for a binary tree node.
class TreeNode:
    def __init__(self, x):
        self.val = x
        self.left = None
        self.right = None

class Solution:
    def minDepth(self, root):
        """
        :type root: TreeNode
        :rtype: int
        """
        if root is None:
            return 0
        nodeList = []
        nodeList.append(root)
        minVal = 1
        while nodeList:
            tmpList = []
            for node in nodeList:
                if node.left is None and node.right is None:
                    return minVal
                if node.left is not None:
                    tmpList.append(node.left)
                if node.right is not None:
                    tmpList.append(node.right)
            minVal += 1
            nodeList = tmpList
        return minVal


#! /usr/bin/env python3
# coding: utf8

"""
File : converter.py
Author : lgbarrere
Brief : Convert a CNF-SAT formula into its ILP version
"""
import os
from os import path, listdir
import shutil
import logging as lg
import time

import pulp

from .utility import Constants, to_ilp_suffix, lines_from_file, split, build_path

class ILPFormula:
    """
    Brief : Defines a local ILP formula
    """
    __PREFIX = 'z' # Prefix of converted variables
    __OBJECTIVE = f'Maximize\n  Obj: {__PREFIX}\n' # Objective

    def __init__(self):
        self.converted = False # Conversion state
        self.constraint_dict = {} # ILP constraints {name : (term_list, goal)}
        self.binary_list = [] # Binaries
        self.nb_variables = 0 # Number of variables
        self.nb_clauses = 0 # Number of clauses

    ## Getters
    def get_prefix(self):
        """
        Brief : Get the prefix of ILP variables
        Return : The prefix
        """
        return self.__PREFIX

    def get_objective(self):
        """
        Brief : Get the Objective function string of ILP variables
        Return : The objective function string
        """
        return self.__OBJECTIVE

    def __repr__(self):
        if self.converted:
            constraint_list = []
            constraint_list.append('Subject To\n')
            for name in self.constraint_dict :
                expr, goal = self.constraint_dict[name]
                expr_str = ' '.join(expr)
                constraint_list.append(f'  {name}: {expr_str} >= {goal}\n')
            tmp_binaries = '\n  '.join(self.binary_list)
            binaries = f'Binary\n  {tmp_binaries}\n'
            return self.__OBJECTIVE + ''.join(constraint_list) + \
                   binaries + 'End'
        lg.warning("No CNF-SAT has been converted into ILP so far.")
        return 'Formula : None\n'

class Converter(Constants):
    """
    Brief : Allows to convert a CNF-SAT into an ILP instance
    with a given file name
    """
    def __init__(self):
        super().__init__()
        self._formula_dict = {} # Formula associated to a file name
        self.__ilp_string_dict = {} # String of ILP associated to a file name
        self.__convert_time_dict = {} # Conversion time associated to a file name

    ## Getters
    def get_ilp_string(self, file_name):
        """
        Brief : Get the string of an ILP from file name
        Return : The ILP string
        > file_name : Name of the DIMACS or ILP file concerned
        """
        return self.__ilp_string_dict[to_ilp_suffix(file_name)]

    def is_converted(self, file_name):
        """
        Brief : Check if the formula has already been converted
        Return : True is a convertion has been done, False otherwise
        > file_name : Name of the DIMACS or ILP file concerned
        """
        return self._formula_dict[to_ilp_suffix(file_name)].converted

    def get_convert_time(self, file_name):
        """
        Brief : Get the time spend to convert the file file_name
        Return : The conversion time if it exists, None otherwise
        > file_name : Name of the DIMACS or ILP file concerned
        """
        file_name = to_ilp_suffix(file_name)
        if file_name in self.__convert_time_dict:
            return self.__convert_time_dict[file_name]
        return None

    def get_constraint(self, file_name):
        """
        Brief : Get the constraints in a dictionary
        Return : The constraints
        > file_name : Name of the DIMACS or ILP file concerned
        """
        return self._formula_dict[to_ilp_suffix(file_name)].constraint_dict

    def get_binaries(self, file_name):
        """
        Brief : Get the binaries in a list
        Return : The binaries
        > file_name : Name of the DIMACS or ILP file concerned
        """
        return self._formula_dict[to_ilp_suffix(file_name)].binary_list

    def get_nb_variables(self, file_name):
        """
        Brief : Get the number of ILP variables
        Return : The number of ILP variables
        > file_name : Name of the DIMACS or ILP file concerned
        """
        return self._formula_dict[to_ilp_suffix(file_name)].nb_variables

    def get_nb_clauses(self, file_name):
        """
        Brief : Get the number of clauses
        Return : The number of clauses
        > file_name : Name of the DIMACS or ILP file concerned
        """
        return self._formula_dict[to_ilp_suffix(file_name)].nb_clauses

    ## Methods
    def print_ilp(self, file_name):
        """
        Brief : Print the formula linked to the given file name
        Return : None
        > file_name : Name of the DIMACS or ILP file concerned
        """
        print(self._formula_dict[to_ilp_suffix(file_name)])

    def __read_dimacs_file(self, file_name, option_folder=None):
        """
        Brief : Get every lines from a dimacs file
        Return : Every lines in a list
        > file_name : The name of the file to read
        > option_folder : If given, search the file in this sub folder
        """
        lg.debug("Reading file %s.", file_name)
        file_path = build_path(
            self.get_data_path(), option_folder, file_name
            )
        return lines_from_file(file_path)

    def __read_ilp_file(self, file_name, option_folder=None):
        """
        Brief : Get every lines from an ILP file
        Return : Every lines in a list
        > file_name : The name of the file to read
        > option_folder : If given, search the file in this sub folder
        """
        lg.debug("Reading file %s.", file_name)
        file_path = build_path(
            self.get_save_path(), option_folder, file_name
            )
        return lines_from_file(file_path)

    def __convert_to_ilp(self, lines, file_name):
        """
        Brief : Convert a CNF-SAT formula into ILP
        Return : None
        > lines : The lines got from the file
        > file_name : The name of the file to get formula
        """
        # Initializations
        file_name = to_ilp_suffix(file_name)
        formula = ILPFormula()
        self._formula_dict[file_name] = formula
        formula.converted = False
        binary_dict = {} # Match SAT variables to ILP variables
        # Consider key 0 as objective
        prefix = formula.get_prefix()
        binary_dict[0] = prefix
        nb_clauses = 0

        # Process
        start_time = time.time()
        for line in lines:
            words = line.split()
            # Ignore DIMACS comments
            if words[0] == 'c' or words[0] == '\n':
                continue
            # Get the number of variables and clauses
            if words[0] == 'p':
                i = 1
                if words[i] == 'cnf':
                    i += 1
                formula.nb_variables = int(words[i])
                formula.nb_clauses = int(words[i + 1])
            # Start writting the constraints
            else:
                word_list = []
                i = 0
                val = int(words[i])
                goal = 1
                nb_clauses += 1
                while val != 0:
                    # Get each variable as int (to know if it's positive)
                    if val < 0:
                        # Respect the rule : not(0) = 1 and not(1) = 0
                        val = -val
                        word_list.append('-')
                        goal -= 1
                    elif val > 0 and i > 0:
                        word_list.append('+')
                    # Set the CNF-SAT variables as binaries
                    if val not in binary_dict :
                        binary_dict[val] = f'{prefix}{val}'
                    word_list.append(binary_dict[val])
                    i += 1
                    val = int(words[i])
                formula.constraint_dict[f'C{nb_clauses}'] = (word_list, goal)
        for binary in binary_dict.values() :
            formula.binary_list.append(binary)
        self.__convert_time_dict[file_name] = time.time() - start_time
        formula.converted = True

    def __ilp_to_string(self, file_name):
        """
        Brief : Convert the ILP from a given file name to a string
        Return : None
        > file_name : The name of the file concerned by  the conversion
        """
        file_name = to_ilp_suffix(file_name)
        formula = self._formula_dict[file_name]
        self.__ilp_string_dict[file_name] = str(formula)

    def convert_from_file(self, file_name, option_folder=None):
        """
        Brief : Convert a CNF-SAT formula into ILP from the given file_name
        Return : None
        > file_name : The name of the file to open
        > option_folder : sub folder with the file to convert
        """
        lines = self.__read_dimacs_file(file_name, option_folder)
        if lines is None:
            return
        file_name = to_ilp_suffix(file_name)
        # If the conversion has not already been set
        if file_name not in self._formula_dict :
            self.__convert_to_ilp(lines, file_name)
            self.__ilp_to_string(file_name)
        else :
            lg.warning("This file has already been converted.")

    def __set_ilp(self, lines, file_name):
        """
        Brief : Set the ILP from a file in the formula dictionary
        Return : None
        > lines : The lines got from the file
        > file_name : The name of the file to open
        """
        # Initializations
        file_name = to_ilp_suffix(file_name)
        formula = ILPFormula()
        self._formula_dict[file_name] = formula
        formula.binary_list = []
        binary_list = []
        # Consider key 0 as objective
        prefix = formula.get_prefix()
        binary_list.append(prefix)
        nb_clauses = 0
        nb_variables = 0
        lines = ''.join(lines)
        lines = split(lines, ['Subject To\n', 'Binary\n', 'End'])

        # Constraints
        for line in lines[0] :
            words = line.split()
            i = 0
            term_list = []
            while words[i] == '' :
                i += 1
            constraint_name = words[i].replace(':', '')
            i += 1
            while words[i] != '>=' :
                term_list.append(words[i])
                i += 1
            goal = int(words[i + 1])
            nb_clauses += 1
            formula.constraint_dict[constraint_name] = (term_list, goal)

        # Binaries
        for line in lines[1] :
            words = line.split()
            i = 0
            while words[i] == '' :
                i += 1
            if words[i] == prefix :
                continue
            if words[i] not in binary_list :
                binary_list.append(words[i])
            nb_variables += 1

        for binary in binary_list :
            formula.binary_list.append(binary)
        formula.nb_variables = nb_variables
        formula.nb_clauses = nb_clauses
        formula.converted = True

    def ilp_from_file(self, file_name, option_folder=None):
        """
        Brief : Load ILP from the given file_name
        Return : None
        > file_name : The name of the file to open
        > option_folder : sub folder with the file to convert
        """
        lines = self.__read_ilp_file(file_name, option_folder)
        if lines is None:
            return
        file_name = to_ilp_suffix(file_name)
        # If the conversion has not already been set
        if file_name not in self._formula_dict :
            self.__set_ilp(lines, file_name)
            self.__ilp_to_string(file_name)
        else :
            lg.warning("This ilp file has already been read.")

    def ilp_from_folder(self, option_folder=None):
        """
        Brief : Load all ILP files from the given folder_name
        Return : None
        > option_folder : The name of the folder to convert files from
        """
        folder_path = build_path(self.get_save_path(), option_folder)
        relative_path = build_path(self.get_save_folder(), option_folder)
        lg.debug("Reading ILP files in folder %s.", relative_path)
        lg.disable(level=lg.DEBUG)
        for file_name in listdir(folder_path):
            file_path = path.join(folder_path, file_name)
            if path.isfile(file_path):
                try:
                    with open(file_path, 'r'):
                        self.ilp_from_file(file_name, option_folder)
                except FileNotFoundError as error:
                    lg.critical("The file was not found. %s", error)
        lg.disable(level=lg.NOTSET)
        lg.debug("Folder read !")

    def convert_from_folder(self, option_folder=None):
        """
        Brief : Convert all CNF files into ILP files from the given folder_name
        Return : None
        > option_folder : The name of the folder to convert files from
        """
        folder_path = build_path(self.get_data_path(), option_folder)
        relative_path = build_path(self.get_data_folder(), option_folder)
        lg.debug("Converting files in folder %s.", relative_path)
        lg.disable(level=lg.DEBUG)
        for file_name in listdir(folder_path):
            file_path = path.join(folder_path, file_name)
            if path.isfile(file_path):
                try:
                    with open(file_path, 'r'):
                        self.convert_from_file(file_name, option_folder)
                except FileNotFoundError as error:
                    lg.critical("The file was not found. %s", error)
        lg.disable(level=lg.NOTSET)
        lg.debug("Folder conversion done !")

    def save_ilp_in_file(self, file_name, option_folder=None):
        """
        Brief : Save a converted ILP in a file (created if missing)
        Return : None
        > file_name : The name of the file to open
        > option_folder : folder to save the file (created if missing)
        """
        if file_name is None :
            lg.critical("No save file name was given.")
            return
        file_name = to_ilp_suffix(file_name)
        # If nothing has been converted earlier, don't save anything
        if not self.is_converted(file_name) :
            lg.warning("Can't save a non converted ILP.")
            return

        lg.debug("Saving in file %s.", file_name)
        folder_path = build_path(self.get_save_path(), option_folder)
        file_path = path.join(folder_path, file_name)
        # Create folder and/or file if missing, then save the ILP
        os.makedirs(folder_path, exist_ok=True)
        if not path.isfile(file_path) :
            with open(file_path, 'w') as file:
                file.write(self.get_ilp_string(file_name))
            lg.debug("Save done !")
        else :
            lg.warning("This save already exists.")

    def save_all_in_folder(self, option_folder=None):
        """
        Brief : Save all converted ILP in a folder
        Return : None
        > option_folder : folder to save the files (created if missing)
        """
        relative_path = build_path(self.get_save_folder(), option_folder)
        lg.debug("Saving all in %s.", relative_path)
        lg.disable(level=lg.DEBUG)
        for file_name in self._formula_dict:
            self.save_ilp_in_file(file_name, option_folder)
        lg.disable(level=lg.NOTSET)
        lg.debug("Folder save done !")

    def clear_save_folder(self, option_folder=None):
        """
        Brief : Delete all files in the given folder_name
        Return : None
        > folder_name : The name of the folder to clear
        """
        folder_path = build_path(self.get_save_path(), option_folder)
        relative_path = build_path(self.get_save_folder(), option_folder)
        lg.debug("Clearing folder %s.", relative_path)
        if not path.exists(folder_path):
            lg.warning("The folder to clear doesn't exist.")
            return
        for file_name in listdir(folder_path):
            file_path = path.join(folder_path, file_name)
            try:
                if path.isfile(file_path):
                    os.remove(file_path)
            except OSError as error:
                lg.critical("The file was not found. %s", error)
        lg.debug("Folder cleared !")

    def clear_all_save_folder(self):
        """
        Brief : Delete everything in the save folder
        Return : None
        """
        folder_path = self.get_save_path()
        lg.debug("Clearing folder %s.", self.get_save_folder())

        for file_name in listdir(folder_path):
            file_path = path.join(folder_path, file_name)
            try:
                if path.isfile(file_path):
                    os.remove(file_path)
                elif path.isdir(file_path):
                    shutil.rmtree(file_path)
            except OSError as error:
                lg.critical("Failed to delete %s. %s", file_name, error)
        lg.debug("Folder cleared !")

class SolverInformation:
    """
    Brief : Give solver information
    """
    def __init__(self, solver=None, status=pulp.LpStatusUndefined):
        self.solver = solver
        self.status = status
        self.time = 0

    def is_solved(self):
        """
        Brief : Check if the problem has been solved or not
        Return : True if the problem is solved, False otherwise
        """
        return self.status != pulp.LpStatusNotSolved and \
               self.status != pulp.LpStatusUndefined

    ## Getters
    def get_status(self):
        """
        Brief : Getter for the status
        Return : The status in a string
        """
        return str(pulp.LpStatus[self.status])

    def get_time(self):
        """
        Brief : Getter for the execution time
        Return : The time
        """
        return self.time

    def __repr__(self):
        return 'Solver : ' + self.solver.name + \
               '\nStatus : ' + self.get_status() + \
               '\nExecution time : ' + str(self.time) + '\n'

class PulpProblem:
    """
    Brief : Defines a PuLP problem
    """
    def __init__(self):
        self.problem = None
        self.var_dict = {}
        self.solver_dict = {}

    ## Getters
    def get_solver_info(self, solver_name='PULP_CBC_CMD'):
        """
        Brief : Getter for the solver information (from its name)
        Return : The solver information
        > solver_name : The name of the solver to get the information
        """
        return self.solver_dict[solver_name]

    def __repr__(self):
        print_list = []
        print_list.append('Problem : ' + str(self.problem))
        print_list.append('Objective : ' + str(self.problem.objective))
        for solver_info in self.solver_dict.values() :
            print_list.append(str(solver_info))
        return '\n'.join(print_list)

class PulpConverter(Converter):
    """
    Brief : Get an ILP instance to create its PuLP version
    """
    def __init__(self):
        super().__init__()
        self.__avail_solver_dict = {}
        for solver_name in pulp.listSolvers(onlyAvailable=True):
            self.__avail_solver_dict[solver_name] = None
        self.__problem_dict = {}

    ## Getters
    def get_problem(self, file_name):
        """
        Brief : Get the ILP problem from a file name
        Return : The ILP problem
        > file_name : The name of the file to get the ILP problem
        """
        return self.__problem_dict[to_ilp_suffix(file_name)]

    def get_solvers(self):
        """
        Brief : Get the list of the available solvers
        Return : The solver list
        """
        return self.__avail_solver_dict.keys()

    def get_solver_path(self, solver_name):
        """
        Brief : Get the path of a solver from its name
        Return : The path in a string (if no path, return None)
        > solver_name : The name of the solver
        """
        return self.__avail_solver_dict[solver_name]

    ## Methods
    def add_solver(self, solver_name, solver_path):
        """
        Brief : Check if a solver can be used with the given path, then add it
        Return : True if the solver is added, False otherwise
        > solver_name : The name of the solver
        > solver_path : The path to the executable of the solver
        """
        model = pulp.LpProblem('Check_solver', pulp.LpMinimize)
        model += pulp.LpVariable('v') == 1
        try:
            model.solve(pulp.getSolver(solver_name, path=solver_path))
            self.__avail_solver_dict[solver_name] = solver_path
            return True
        except pulp.apis.core.PulpSolverError:
            lg.critical(
                "The solver requested " + solver_name + \
                " can\'t be used with the path " + solver_path
                )
        return False

    def define_problem(self, file_name, name='NoName'):
        """
        Brief : Define the problem from a loaded ILP given by a file name
        Return : True if the problem is defined, False otherwise
        > file_name : The name of the ILP file
        > name : An optional name for the problem
        """
        file_name = to_ilp_suffix(file_name)
        if not self.is_converted(file_name) :
            return False
        if file_name in self.__problem_dict :
            lg.warning("Problem already defined.")
            # If a new solver has been registered, add it to this problem
            pulp_problem = self.__problem_dict[file_name]
            for solver_name in self.__avail_solver_dict.keys():
                if solver_name not in pulp_problem.solver_dict:
                    pulp_problem.solver_dict[solver_name] = SolverInformation(
                        pulp.getSolver(solver_name), pulp.LpStatusNotSolved
                        )
            return False

        pulp_problem = PulpProblem()
        self.__problem_dict[file_name] = pulp_problem
        pulp_problem.problem = pulp.LpProblem(name=name, sense=pulp.LpMaximize)

        # Variables
        binary_list = self.get_binaries(file_name)
        pulp_problem.var_dict = {}
        for value in binary_list:
            pulp_problem.var_dict[value] = pulp.LpVariable(
                name=value, cat=pulp.LpBinary
                )

        # Constraints
        constraint_dict = self.get_constraint(file_name)
        for constraint in constraint_dict.values() :
            var_list = []
            term = constraint[0]
            for (i, _) in enumerate(term) :
                sign = 1
                if term[i] == '-':
                    sign = -1
                    i += 1
                elif term[i] == '+':
                    i += 1
                var_list.insert(0, (pulp_problem.var_dict[term[i]], sign))
                i += 1

            pulp_problem.problem += pulp.LpConstraint(
                e=pulp.LpAffineExpression(e=var_list),
                sense=pulp.LpConstraintGE,
                rhs=constraint[1]
                )

        # Objective function
        pulp_problem.problem += pulp.LpVariable(
            name=self._formula_dict[file_name].get_prefix(),
            lowBound=1, upBound=1
            )

        for solver_name in self.__avail_solver_dict.keys() :
            pulp_problem.solver_dict[solver_name] = SolverInformation(
                pulp.getSolver(solver_name), pulp.LpStatusNotSolved
                )
        return True

    def define_problem_from_folder(self, option_folder=None, name='NoName'):
        """
        Brief : Define problems from a folder converted into ILP,
        if the given folder is None, problems are defined from the saves folder
        Return : None
        > option_folder : The name of the sub folder to define problems from
        > name : The name of the problems
        """
        folder_path = build_path(self.get_save_path(), option_folder)
        relative_path = build_path(self.get_save_folder(), option_folder)
        lg.debug("Define all from %s.", relative_path)
        lg.disable(level=lg.DEBUG)
        for file_name in listdir(folder_path) :
            file_path = path.join(folder_path, file_name)
            if path.isfile(file_path) :
                self.define_problem(file_name, name)
        lg.disable(level=lg.NOTSET)
        lg.debug("Folder definition done !")

    def solve(self, file_name, solver_name='PULP_CBC_CMD', path=None, time_limit=None):
        """
        Brief : Solve a problem from its file name, updating its information
        Return : None
        > file_name : The name of the file used to define a problem to solve
        > solver_name : The name of the solver to use
        > time_limit : The maximum time taken to solve before interruption
        """
        file_name = to_ilp_suffix(file_name)
        pulp_problem = self.__problem_dict[file_name]
        solver_info = pulp_problem.solver_dict[solver_name]
        if not solver_info.is_solved() :
            # If the solver is interrupted, consider the problem is unsolved
            try :
                start_time = time.time()
                solver_info.status = pulp_problem.problem.solve(
                    pulp.getSolver(solver_name, path=path, timeLimit=time_limit)
                    )
                if solver_info.status == pulp.LpStatusNotSolved:
                    solver_info.time = 'Timeout'
                else:
                    solver_info.time = time.time() - start_time
            except pulp.apis.core.PulpSolverError :
                lg.warning("Interrupted solver.")
                solver_info.status = pulp.LpStatusNotSolved
                solver_info.time = 'Stopped'
        else :
            lg.warning("Undefined problem or already solved.")

    def solve_folder(
        self, option_folder=None, solver_name='PULP_CBC_CMD', time_limit=None
        ):
        """
        Brief : Start solving an entire folder and set the time to proceed
        Return : None
        > option_folder : The name of the sub folder to define problems from
        > solver_name : The name of the solver to use
        > time_limit : The maximum time taken to solve before interruption
        """
        folder_path = build_path(self.get_save_path(), option_folder)
        relative_path = build_path(self.get_save_folder(), option_folder)
        lg.debug("Solving all from %s.", relative_path)
        lg.disable(level=lg.DEBUG)
        for file_name in listdir(folder_path) :
            file_path = path.join(folder_path, file_name)
            if path.isfile(file_path) :
                self.solve(file_name, solver_name, time_limit)
        lg.disable(level=lg.NOTSET)
        lg.debug("Folder solve done !")

    def save_results(self, result_file):
        """
        Brief : Save all solved ILP results in a file (created if missing)
        Return : None
        """
        # If the file has been solved
        folder = self.get_result_folder()
        lg.debug("Saving solutions in folder %s.", folder)
        folder_path = path.join(self.get_root_path(), folder)
        file_path = path.join(folder_path, result_file)
        # Create folder and/or file if missing, then save the solutions
        os.makedirs(folder_path, exist_ok=True)
        with open(file_path, 'a') as file:
            for (file_name, problem) in self.__problem_dict.items() :
                file.write(f'File : {file_name}\n')
                convert_time = self.get_convert_time(file_name)
                file.write(f'  Conversion time : {convert_time}\n')
                for solver_name in problem.solver_dict :
                    solver_info = problem.solver_dict[solver_name]
                    if solver_info.status != pulp.LpStatusNotSolved \
                       or solver_info.time != 0:
                        file.write(f'  Solver : {solver_name}')
                        file.write(f' | Status : {solver_info.get_status()}')
                        file.write(f' | Execution time : {solver_info.time}\n')
        lg.debug("Solutions Saved !")

def main():
    """
    Brief : Test some uses of this library
    Return : None
    """
    # Verbose debug and prepare the tests
    lg.basicConfig(level=lg.DEBUG)
    test_folder_name = 'dimacs'
    converter = PulpConverter()
    converter.clear_all_save_folder()

    # Test limits in example files
    test_file = 'test.cnf'
    converter.convert_from_file(test_file)
    converter.save_ilp_in_file(test_file)
    print('n = ' + str(converter.get_nb_variables(test_file)))
    print('m = ' + str(converter.get_nb_clauses(test_file)))
    converter.convert_from_folder()

    # Test to get an ILP directly from an ILP file
    ilp_test = to_ilp_suffix(test_file)
    converter.ilp_from_file(ilp_test)
    converter.print_ilp(ilp_test)

    # Test with PuLP definitions
    lg.disable(level=lg.DEBUG)
    example_file = 'example.cnf'
    converter.convert_from_file(example_file)
    converter.define_problem(example_file)
    converter.solve(example_file)
    converter.save_results()

    # Test with important data folder
    lg.disable(level=lg.NOTSET)
    converter.convert_from_folder(test_folder_name)
    converter.save_all_in_folder(test_folder_name)

if __name__ == "__main__":
    main()


import Adafruit_DHT
import logging
import threading

from kcam import observer

LOG = logging.getLogger(__name__)

devices = {
    'DHT22': Adafruit_DHT.DHT22,
    'DHT11': Adafruit_DHT.DHT11,
}

class TemperatureSensor(threading.Thread):
    default_interval = 60
    default_devicetype = Adafruit_DHT.DHT22

    def __init__(self, pin,
                 devicetype=None,
                 interval=None,
                 **kwargs):
        super(TemperatureSensor, self).__init__(**kwargs)
        self.pin = pin
        self.interval = interval if interval else self.default_interval
        self.devicetype = (devices[devicetype] if devicetype
                           else self.default_devicetype)
        self.evt_stop = threading.Event()

        self.temperature = observer.Value()
        self.humidity = observer.Value()
        self.mutex = threading.RLock()

    def __str__(self):
        return '<TemperatureSensor pin:%d interval:%d>' % (
            self.pin, self.interval)

    def stop(self):
        self.evt_stop.set()

    def run(self):
        LOG.info('starting temperature sensor %s', self)

        while True:
            humidity, temperature = Adafruit_DHT.read_retry(
                self.devicetype, self.pin)
            if humidity is not None and temperature is not None:
                LOG.info('got humidity = %f, temp = %f',
                         humidity, temperature)
                with self.mutex:
                    self.temperature.set(temperature)
                    self.humidity.set(humidity)

            else:
                LOG.error('failed to read temperature sensor')

            if self.evt_stop.wait(self.interval):
                break

        LOG.info('stopping temperature sensor on pin %d', self.pin)

    @property
    def value(self):
        with self.mutex:
            return dict(temperature=self.temperature.value,
                        humidity=self.humidity.value)


from typing import List, Literal

from fastapi import APIRouter, Depends, HTTPException, status
from pydantic import BaseModel

from app.infrastructure.base import COMMAND, get_db
from app.infrastructure.now import get_now
from app.infrastructure.redis_producer import RedisProducer, get_queue, get_rabbit_queue
from app.infrastructure.video_repository import VideoData, VideoRepository
from app.interface.router.auth import get_current_user

router = APIRouter(prefix="/video", tags=["video"])

@router.get("/")
async def read_all_videos(dynamodb=Depends(get_db)):
    repo = VideoRepository(dynamodb=dynamodb)
    return repo.get_all()["Items"]

@router.get("/{key}")
async def read_one_video(key: str, dynamodb=Depends(get_db)):
    repo = VideoRepository(dynamodb=dynamodb)
    return repo.get(key)

@router.delete("/{key}")
async def delete_video(
    key: str, dynamodb=Depends(get_db), username=Depends(get_current_user)
):
    repo = VideoRepository(dynamodb=dynamodb)
    res = repo.delete(key=key, username=username)
    if res:
        return {"deleted": key}
    raise HTTPException(
        status_code=status.HTTP_404_NOT_FOUND, detail=f"there is no element with {key}"
    )

class UploadVideo(BaseModel):
    key: str

@router.post("/upload")
async def upload_video(
    request: UploadVideo, dynamodb=Depends(get_db), username=Depends(get_current_user)
):
    repo = VideoRepository(dynamodb=dynamodb)
    current_time = get_now()
    ret = repo.put(
        user_name=username["user"], current_time=current_time, key=request.key
    )
    if ret:
        return {"uploaded": request.key}
    raise HTTPException(
        status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"db error"
    )

class ExtractFeature(COMMAND):
    type: Literal["ExtractFeature"]
    key: str

@router.post("/extract")
async def extract_feature(
    request: ExtractFeature,
    dynamodb=Depends(get_db),
    username=Depends(get_current_user),
    producer=Depends(get_rabbit_queue),
):
    repo = VideoRepository(dynamodb=dynamodb)
    ret = repo.get(key=request.key)
    if ret is None:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail=f"no such video of {request.key} key",
        )
    ret = repo.put(
        user_name=ret.user_name,
        current_time=ret.start_time,
        key=ret.key,
        status="in_queue"
    )
    producer.put(request)

class ShortenVideo(COMMAND):
    type: Literal["ShortenVideo"]
    key: str
    must_include_feature: List[str]

@router.post("/shorten")
async def shorten_video(
    request: ShortenVideo,
    dynamodb=Depends(get_db),
    producer=Depends(get_rabbit_queue),
):
    repo = VideoRepository(dynamodb=dynamodb)
    ret = repo.get(key=request.key)
    if ret is None:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail=f"no such video of {request.key} key",
        )
    ret = repo.put(
        user_name=ret.user_name,
        current_time=ret.start_time,
        key=ret.key,
        status="shorten_start"
    )
    producer.put(request)


from tkinter import*

WIDTH = 500
HEIGHT = 500

class Labirint():
    def __init__(self, master):

        self.canvas = Canvas(master, width=WIDTH, height=HEIGHT)
        self.canvas.pack()
        self.nalozi_matriko()
        zac_x, zac_y = self.narisi()
        self.k = Krogla(zac_x, zac_y, self.canvas, self.matrika)
        self.canvas.bind("<Left>", self.k.levo)
        self.canvas.bind("<Right>", self.k.desno)
        self.canvas.bind("<Up>", self.k.gor)
        self.canvas.bind("<Down>", self.k.dol)
        self.canvas.focus_set()
        menu = Menu(master)
        master.config(menu=menu)
        file = Menu(menu)
        menu.add_cascade(label="Menu", menu = file)
        file.add_command(label="Shrani", command=self.shrani)
        file.add_command(label="Zapri", command=master.destroy)

    def nalozi_matriko(self, file="nacrt2.txt"):
        self.matrika = []
        with open(file) as beri:
            for vrstica in beri:
                podsez = []
                for znak in vrstica:
                    if znak != "\n":
                        podsez.append(znak)
                self.matrika.append(podsez)

    def narisi(self):
        for i in range(len(self.matrika)):
            for j in range(len(self.matrika[i])):
                if self.matrika[i][j] == "*":
                    self.canvas.create_rectangle(
                        50*j, 50*i, 50*j + 50, 50*i + 50, fill = "black")
                if self.matrika[i][j] == "M":
                    self.canvas.create_rectangle(
                        50*j, 50*i, 50*j + 50, 50*i + 50, fill = "yellow")
                if self.matrika[i][j] == "Z":
                    zac_x = int(j * 50)
                    zac_y = int(i * 50)
        return zac_x, zac_y

    def shrani(self):
        print ("SHRANJENO")
        x = self.k.x / 50
        y = self.k.y / 50
        with open ("nacrt2.txt", "wt") as pisi:
            for i in range (len(self.matrika)):
                for j in range(len(self.matrika[i])):
                    if self.matrika[i][j] == "Z":
                        print(" ",file=pisi,end="")
                    elif i == y and j == x:
                        print("Z", file=pisi,end="")
                    else:
                        print(self.matrika[i][j],file=pisi,end="")
                print("",file=pisi)

class Krogla():
    def __init__(self, x, y, canvas, zid):
        self.zid = zid
        self.canvas = canvas
        self.x = x
        self.y = y
        self.krogec = self.canvas.create_oval(
            self.x, self.y, self.x + 50, self.y + 50, fill="orange")

    def levo(self, event):
        new_x = int((self.x - 50) / 50)
        new_y = int(self.y / 50)
        index = True

        if index:
            if self.zid[new_y][new_x] != "*" and new_x >=0:
                self.x -= 50
                self.canvas.coords(self.krogec, self.x, self.y, self.x + 50, self.y + 50)
                self.preveri_ce_je_konec(new_x, new_y)

    def desno(self, event):
        new_x = int((self.x + 50) / 50)
        new_y = int(self.y / 50)
        index = True
        try:
            self.zid[new_y][new_x]
        except IndexError:
            index = False

        if index:
            if self.zid[new_y][new_x] != "*":
                self.x += 50
                self.canvas.coords(self.krogec, self.x, self.y, self.x + 50, self.y + 50)
                self.preveri_ce_je_konec(new_x, new_y)

    def gor(self, event):
        new_x = int(self.x  / 50)
        new_y = int((self.y - 50) / 50)
        index = True

        if index:
            if self.zid[new_y][new_x] != "*" and new_y >= 0:
                self.y -= 50
                self.canvas.coords(self.krogec, self.x, self.y, self.x + 50, self.y + 50)
                self.preveri_ce_je_konec(new_x, new_y)

    def dol(self, event):
        new_x = int(self.x  / 50)
        new_y = int((self.y + 50) / 50)
        index = True
        try:
            self.zid[new_y][new_x]
        except IndexError:
            index = False

        if index:
            if self.zid[new_y][new_x] != "*":
                self.y += 50
                self.canvas.coords(self.krogec, self.x, self.y, self.x + 50, self.y + 50)
                self.preveri_ce_je_konec(new_x, new_y)

    def preveri_ce_je_konec(self,x, y):
        if self.zid[y][x] == "M":
            print ("ZMAGA")

root = Tk()

aplikacija = Labirint(root)

root.mainloop()


# Create a program that reads the year of birth of seven people. In the end, show how many of those people are over 18
# And how many aren't

from datetime import date

over_18 = 0
under_18 = 0

for i in range(1, 8):
    yob = int(input('Year of birth {}: '.format(i)))
    age = date.today().year - yob
    if age >= 18:
        over_18 += 1
    else:
        under_18 += 1

print("""There are \033[1;34m{}\033[m people \033[1;34m18 or over\033[m
and also \033[1;32m{}\033[m people \033[1;32munder 18\033[m.""".format(over_18, under_18))


#!/bin/python3

import math
import os
import random
import re
import sys

#
# Complete the 'repeatedString' function below.
#
# The function is expected to return a LONG_INTEGER.
# The function accepts following parameters:
#  1. STRING s
#  2. LONG_INTEGER n
#

"""
There is a string, s, of lowercase English letters that is repeated infinitely many times.
Given an integer, n, find and print the number of letter a's in the first n letters of the infinite string.
"""

def repeatedString(s, n):
    # Write your code here
    numA = 0
    # takes care of case where number of chars to check is less than length of s
    if len(s)>n:
        for i in range(n):
            if s[i] == 'a':
                numA+=1
        return numA
    # below takes care of every other case
    # count number of s repeats
    numRepeats = n//len(s)
    # count number of chars in the last incomplete string if any
    numRemainder = n%len(s)
    # count number of 'a' in this incomplete string
    numAextra = 0
    # count number of 'a' in s
    for i in range(len(s)):
        if s[i] == 'a':
            numA+=1
            if i<numRemainder:
                numAextra+=1
    return numA*max(numRepeats, 1)+numAextra

if __name__ == '__main__':

    s = input()

    n = int(input().strip())

    result = repeatedString(s, n)

    print(result)


"""
Get the surprisals of each sentence based on a specific language model.
This script was partly adapted from code in the repo:
    https://github.com/krishnap25/mauve
"""

import sys
import os
import argparse
from tqdm import tqdm
import torch
from torch.functional import F
from transformers import GPT2LMHeadModel, AutoTokenizer

sys.path.insert(1, os.path.join(sys.path[0], '..'))
from util import constants
from util import util

def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--text-fpath", help="The path to save text data",
                        type=str, required=True)
    parser.add_argument("--surprisals-fpath", help="The path to save text data",
                        type=str, required=True)
    parser.add_argument("--model", help="The model to use to get surprisals",
                        type=str, required=True)
    parser.add_argument("--batch-size", type=int, default=4)
    parser.add_argument("--dump-size", type=int, default=10000)
    args = parser.parse_args()
    print(args)

    return args

def get_model(model_name, tokenizer):
    if 'gpt2' in model_name:
        model = GPT2LMHeadModel.from_pretrained(
            model_name, pad_token_id=tokenizer.eos_token_id)
        model = model.to(constants.device)
        model = model.eval()
    else:
        raise ValueError(f'Unknown model: {model_name}')
    return model

def get_tokenizer(model_name='gpt2'):
    if 'gpt2' in model_name:
        tokenizer = AutoTokenizer.from_pretrained(model_name)
    else:
        raise ValueError(f'Unknown model: {model_name}')
    return tokenizer

def count_parameters(model):
    return sum(x.numel() for x in model.parameters() if x.requires_grad)

def get_batch_surprisals(model, batch):
    padded_batch = torch.nn.utils.rnn.pad_sequence(
        batch, batch_first=True, padding_value=0).to(constants.device)
    attention_mask = torch.nn.utils.rnn.pad_sequence(
        [torch.ones(sent.size(0)) for sent in batch],
        batch_first=True, padding_value=0).long().to(constants.device)

    outs = model(input_ids=padded_batch, attention_mask=attention_mask,
                 past_key_values=None, return_dict=True)

    # Get surprisals of following word pieces
    surprisals_full = - F.log_softmax(outs.logits, dim=-1)
    surprisals_sentence = torch.gather(
        surprisals_full[:, :-1], -1, padded_batch[:, 1:].unsqueeze(-1)).squeeze(-1)
    surprisals_sentence = (surprisals_sentence * attention_mask[:, 1:])

    return surprisals_sentence.sum(-1).cpu()

@torch.no_grad()
def get_surprisals_from_tokenised_text(model, tokenized_texts, batch_size):
    surprisals = []

    with tqdm(total=len(tokenized_texts), desc="Getting chunk surprisals") as pbar:
        for batch in util.chunker(tokenized_texts, batch_size):
            surprisals_batch = get_batch_surprisals(model, batch)

            surprisals.append(surprisals_batch)
            pbar.update(len(batch))

    return torch.cat(surprisals).numpy()

def get_surprisals_from_input(texts, model_name, batch_size, max_len=1024):
    tokenizer = get_tokenizer(model_name)
    tokenized_texts = [
        tokenizer.encode(
            tokenizer.bos_token + sen + tokenizer.eos_token,
            return_tensors='pt', truncation=True, max_length=max_len).view(-1)
        for sen in texts
    ]

    model = get_model(model_name, tokenizer)
    # print('Model %s has %dM parameters' % (model_name, count_parameters(model) * 1e-6))

    surprisals = get_surprisals_from_tokenised_text(
        model, tokenized_texts, batch_size=batch_size)

    return surprisals

def save_surprisals(ids, surprisals, fpath):
    data = {
        'ids': ids,
        'surprisals': surprisals,
    }
    util.write_data(fpath, data)

def load_surprisals(fpath):
    data = util.read_data(fpath)
    return data['ids'], data['surprisals']

def try_load_surprisals(fpath):
    if os.path.isfile(fpath):
        return load_surprisals(fpath)
    return None, None

def process_data(text_fpath, surprisals_fpath, model, batch_size, dump_size):
    text_data = util.read_data(text_fpath)
    ids = text_data['ids']
    texts = text_data['texts']
    model = constants.MODEL_NAMES[model]

    # If dump_size is -1 save all results only once, else save temporary results
    if dump_size != -1:
        surprisals_fpath_temp = surprisals_fpath + '.temp'
        surprisals = util.process_and_save_temp(
            ids, texts, function=get_surprisals_from_input, dump_size=dump_size,
            function_kwargs={'model_name': model, 'batch_size': batch_size, 'max_len': 1024},
            load_function=load_surprisals, save_function=save_surprisals,
            temp_fpath=surprisals_fpath_temp)
    else:
        surprisals = get_surprisals_from_input(
            texts=texts, model_name=model, batch_size=batch_size, max_len=1024)

    save_surprisals(ids, surprisals, surprisals_fpath)

def main():
    args = get_args()
    print('Is PyTorch in cuda? ', torch.cuda.is_available())
    print('Is PyTorch device: ', constants.device)

    process_data(args.text_fpath, args.surprisals_fpath,
                 args.model, batch_size=args.batch_size, dump_size=args.dump_size)

if __name__ == '__main__':
    main()


import logging
import re
from itertools import zip_longest

from tuxemon.compat.rect import Rect
from tuxemon.core import prepare
from tuxemon.core.locale import T

logger = logging.getLogger(__name__)

def get_cell_coordinates(rect, point, size):
    """Find the cell of size, within rect, that point occupies."""
    cell = [None, None]
    point = (point[0] - rect.x, point[1] - rect.y)
    cell[0] = (point[0] // size[0]) * size[0]
    cell[1] = (point[1] // size[1]) * size[1]
    return tuple(cell)

def transform_resource_filename(*filename):
    """ Appends the resource folder name to a filename

    :param filename: String
    :rtype: basestring
    """
    return prepare.fetch(*filename)

def new_scaled_rect(*args, **kwargs):
    """ Create a new rect and scale it

    :param args: Normal args for a Rect
    :param kwargs: Normal kwargs for a Rect
    :rtype: tuxemon.compat.rect.Rect
    """
    rect = Rect(*args, **kwargs)
    return scale_rect(rect)

def scale_rect(rect, factor=prepare.SCALE):
    """ Scale a rect.  Returns a new object.

    :param rect: Rect
    :param factor: int
    :rtype: tuxemon.compat.rect.Rect
    """
    return Rect([i * factor for i in list(rect)])

def scale_sequence(sequence):
    """ Scale the thing

    :param sequence:
    :rtype: list
    """
    return [i * prepare.SCALE for i in sequence]

def scale(number):
    """ Scale the thing

    :param number: int
    :rtype: int
    """
    return prepare.SCALE * number

def check_parameters(parameters, required=0, exit=True):
    """
    Checks to see if a given list has the required number of items
    """
    if len(parameters) < required:
        import inspect
        calling_function = inspect.stack()[1][3]
        logger.error("'" + calling_function + "' requires at least " + str(required) + "parameters.")
        if exit:
            import sys
            sys.exit(1)
        return False

    else:
        return True

def calc_dialog_rect(screen_rect):
    """ Return a rect that is the area for a dialog box on the screen

    :param screen_rect:
    :return:
    """
    rect = screen_rect.copy()
    if prepare.CONFIG.large_gui:
        rect.height *= .4
        rect.bottomleft = screen_rect.bottomleft
    else:
        rect.height *= .25
        rect.width *= .8
        rect.center = screen_rect.centerx, screen_rect.bottom - rect.height
    return rect

def open_dialog(session, text, avatar=None, menu=None):
    """ Open a dialog with the standard window size

    :param tuxemon.core.session.Session session: Game session
    :param text: list of strings
    :param avatar: optional avatar sprite
    :param menu: optional menu object

    :rtype: tuxemon.core.states.dialog.DialogState
    """
    rect = calc_dialog_rect(session.client.screen.get_rect())
    return session.client.push_state("DialogState", text=text, avatar=avatar, rect=rect, menu=menu)

def nearest(l):
    """ Use rounding to find nearest tile

    :param l:
    :return:
    """
    return tuple(int(round(i)) for i in l)

def trunc(l):
    return tuple(int(i) for i in l)

def number_or_variable(session, value):
    """ Returns a numeric game variable by its name
    If value is already a number, convert from string to float and return that

    :param session:
    :param value: Union[str, float, int]

    :rtype: float

    :raises: ValueError
    """
    player = session.player
    if value.isdigit():
        return float(value)
    else:
        try:
            return float(player.game_variables[value])
        except (KeyError, ValueError, TypeError):
            logger.error("invalid number or game variable {}".format(value))
            raise ValueError

def cast_values(parameters, valid_parameters):
    """ Change all the string values to the expected type

    This will also check and enforce the correct parameters for actions

    :param parameters:
    :param valid_parameters:
    :return:
    """

    # TODO: stability/testing
    def cast(i):
        ve = False
        t, v = i
        try:
            for tt in t[0]:
                if tt is None:
                    return None

                try:
                    return tt(v)
                except ValueError:
                    ve = True

        except TypeError:
            if v is None:
                return None

            if v == '':
                return None

            return t[0](v)

        if ve:
            raise ValueError

    try:
        return list(map(cast, zip_longest(valid_parameters, parameters)))
    except ValueError:
        logger.error("Invalid parameters passed:")
        logger.error("expected: {}".format(valid_parameters))
        logger.error("got: {}".format(parameters))
        raise

def show_item_result_as_dialog(session, item, result):
    """ Show generic dialog if item was used or not

    :param tuxemon.core.session.Session session: Session
    :param tuxemon.core.item.item.Item item: Item
    :param dict result:
    :return: None
    """
    msg_type = 'use_success' if result['success'] else 'use_failure'
    template = getattr(item, msg_type)
    if template:
        message = T.translate(template)
        open_dialog(session, [message])

def split_escaped(string_to_split, delim=","):
    """Splits a string by the specified deliminator excluding escaped
    deliminators.

    :param string_to_split: The string to split.
    :param delim: The deliminator to split the string by.

    :type string_to_split: Str
    :type delim: Str

    :rtype: List
    :returns: A list of the splitted string.

    """
    # Split by "," unless it is escaped by a "\"
    split_list = re.split(r'(?<!\\)' + delim, string_to_split)

    # Remove the escape character from the split list
    split_list = [w.replace(r'\,', ',') for w in split_list]

    # strip whitespace around each
    split_list = [i.strip() for i in split_list]

    return split_list

def round_to_divisible(x, base=16):
    """ Rounds a number to a divisible base.

    This is used to round collision areas that aren't defined well. This
    function assists in making sure collisions work if the map creator didn't
    set the collision areas to round numbers.

    **Examples:**
    >>> round_to_divisible(31.23, base=16)
    32
    >>> round_to_divisible(17.8, base=16)
    16
    :param x: The number we want to round.
    :param base: The base that we want our number to be divisible by. (Default: 16)
    :type x: Float
    :type base: Integer
    :rtype: Integer
    :returns: Rounded number that is divisible by "base".
    """
    return int(base * round(float(x) / base))

def copy_dict_with_keys(source, keys):
    """ Return new dict using only the keys/value from `keys`

    If key from keys is not present no error is raised

    :param Dict source:
    :param List[str] keys:
    :rtype: Dict
    """
    return {k: source[k] for k in keys if k in source}


import numpy as np
import json
import logging
import os
import ast
from data.datasets import *
from models.spatiotemporal import *
from models.temporal import *
from trainer import Trainer, Mixed, OnlineTrainer, Incrementer

class Configuration(dict):
    """ Dictionary subclass whose entries can be accessed by attributes
        (as well as normally).
    """

    def __init__(self, *args, **kwargs):
        super(Configuration, self).__init__(*args, **kwargs)
        self.__dict__ = self

    @staticmethod
    def from_nested_dict(data):
        """Construct nested AttrDicts from nested dictionaries or from a JSON file.

        Arguments:
            data {dict, JSON} -- Either a dict or JSON path name.

        Returns:
            Configuration -- A Configuration object.
        """
        if isinstance(data, str) and ".json" in data:
            f = open(data)
            data = json.load(f)
            f.close()

            return Configuration({key: Configuration.from_nested_dict(data[key]) for key in data})
        elif not isinstance(data, dict):
            return data
        else:
            return Configuration({key: Configuration.from_nested_dict(data[key]) for key in data})

def set_logger(log_path):
    """Set the logger to log info in terminal and file `log_path`.
        In general, it is useful to have a logger so that every output to the terminal is saved
        in a permanent file. Here we save it to `model_dir/train.log`.
        Example:
        ```
        logging.info("Starting training...")
        ```

    Arguments:
        log_path {[type]} -- The destination directory where to log.
    """
    logger = logging.getLogger()
    logger.setLevel(logging.INFO)

    if not logger.handlers:
        # Logging to a file
        file_handler = logging.FileHandler(log_path, mode='w')
        file_handler.setFormatter(logging.Formatter('%(asctime)s:%(levelname)s: %(message)s'))
        logger.addHandler(file_handler)

        # Logging to console
        stream_handler = logging.StreamHandler()
        stream_handler.setFormatter(logging.Formatter('%(message)s'))
        logger.addHandler(stream_handler)

def str_to_bool(s):
    """Converts a string with a boolean value to a bool.

    Arguments:
        s {str} -- A string of boolean value.

    Raises:
        ValueError: Raised if the string value is not "True" or "False".

    Returns:
        bool -- Returns the value of the boolean string.
    """
    if s == 'True':
         return True
    elif s == 'False':
         return False
    else:
         raise ValueError

def minmax(x, minimum, maximum):
    """A function to for min-max normalization, such that the output is constrained
    in the interval [0,1].

    Arguments:
        x {np.array} -- A numpy array object with the data.
        minimum {double} -- Minimum of the data
        maximum {double} -- Maximum of the data

    Returns:
        np.array -- Min-max normalized data
    """

    y = (x-minimum)/(maximum-minimum)

    return y

def inverse_minmax(y, minimum, maximum):
    """A function to retrieve the original data normalized by min-max.

    Arguments:
        y {np.array} -- Min-max normalized data
        minimum {double} -- Minimum of the data
        maximum {double} -- Maximum of the data

    Returns:
        np.array -- Original numpy array with the data
    """

    x = y * (maximum-minimum) + minimum

    return x

def get_trainer(trainer_name, model, config, criterion, optimizer, dataset, model_dir):
    """A utility function for getting a Trainer object by name. Returns a Trainer object used to
    wrap the training of neural networks.

    Arguments:
        trainer_name {string} -- Name of the Trainer object
        model {nn.Module} -- A PyTorch module with parameters to be optimized
        criterion {nn.loss} -- Loss function to optimize with respect to
        optimizer {nn.optimizer} -- Optimizer object
        dataset {torch.utils.data.Dataset} -- A dataset object to use with dataloaders
        model_dir {str} -- Directory path to the model

    Raises:
        NotImplementedError: Thrown if the Trainer name is incorrect or the Trainer object does not exist.

    Returns:
        Trainer -- Trainer object
    """

    if trainer_name == "trainer":
        trainer = Trainer(model, config, criterion, optimizer, dataset, model_dir)
    elif trainer_name == "approximator" or trainer_name == "online":
        trainer = OnlineTrainer(model, config, criterion, optimizer, dataset, model_dir)
    elif trainer_name == "mixed":
        trainer = Mixed(model, config, criterion, optimizer, dataset, model_dir)
    elif trainer_name == "incrementer":
        trainer = Incrementer(model, config, criterion, optimizer, dataset, model_dir)
    else:
        raise NotImplementedError("Trainer not available.")

    return trainer

def get_optimizer(model, optimizer_name, params):
    """A utility function for getting an optimizer for PyTorch neural networks.

    Arguments:
        model {nn.Module} -- A PyTorch module with parameters to be optimized
        optimizer_name {String} -- The name of the optimizer, either "adam" or "sgd"
        params {Configuration} -- A Configuration object with parameter attributes

    Returns:
        torch.optim -- A PyTorch optimizer object
    """

    if optimizer_name == "adam":
        optimizer   = torch.optim.Adam(model.parameters(), lr=params.lr, betas=(0.9, 0.999), eps=1e-08, weight_decay=params.decay, amsgrad=False)
    elif optimizer_name == "sgd":
        optimizer   = torch.optim.SGD(model.parameters(), lr=params.lr, momentum=params.momentum, nesterov=True, weight_decay=params.decay)
    elif optimizer_name == "rmsprop":
        optimizer   = torch.optim.RMSprop(model.parameters(), lr=params.lr, alpha=0.99, eps=1e-08, weight_decay=params.decay, momentum=0, centered=False)
    else:
        raise NotImplementedError("Optimizer not available.")

    return optimizer

def get_dataset(source, data_path, directory, mode, use_transform, params):
    """Utility function for getting the desired FRAP dataset. Divided into "temporal",
    "spatiotemporal" and "generate", they generate train-test split PyTorch datasets and a
    generator of unique training data from Matlab, respectively.

    Arguments:
        source {str} -- The source of the data, either
                        - "temporal" returns a recovery curve dataset from a directory
                        - "spatiotemporal" returns a FRAP pixel dataset from a directory
                        - "generate"
                        - "mixed"
                        - "files_single" returns a dataset of given size as a single binary file.
                        - "files_split" returns a dataset of given size as one binary file for every sample. Suitable for big data.

        data_path {str} -- Path to the dataset samples
        directory {str} -- Path to the target directory of the generator
        mode {str} -- Either "temporal", "spatiotemporal", "all" or "Fourier", the type of data generated.
        params {Configuration} -- A Configuration dict of parameters, such as batch size.

    Returns:
        torch.nn.utils.Dataset -- A dataset to be loaded to a PyTorch Dataloader.
    """

    if source == "temporal":

        dataset = TemporalDataset(directory = data_path)

    elif source == "spatiotemporal":

        dataset = SpatiotemporalDataset(directory= data_path)

    elif source == "generate":

        dataset = MatlabGenerator(batch_size=params.batch_size,
                                  directory=directory,
                                  mode=mode,
                                  transform=use_transform,
                                  noise_level=params.noise_level,
                                  n_workers=params.batch_size)

    elif source == "mixed":

        dataset = MixedMatlabGenerator(dataset_size = params.dataset_size,
                                       batch_size=params.batch_size,
                                        directory=directory,
                                        mode=mode,
                                        transform=use_transform,
                                        noise_level=params.noise_level,
                                        n_workers=params.batch_size)

    elif source == "fromfiles" or source == "files_single":
        shape = ast.literal_eval(params.shape)

        dataset = FileMatlabGenerator(dataset_size = params.dataset_size,
                                        directory=directory,
                                        transform=use_transform,
                                        noise_level=params.noise_level,
                                        n_workers=params.batch_size,
                                        shape=shape
                                        )

    elif source == "mixed_separate" or source == "files_split":
            shape = ast.literal_eval(params.shape)

            training = FileMatlabGenerator(dataset_size = params.dataset_size,
                                        directory=os.path.join(directory, "train"),
                                        transform=use_transform,
                                        noise_level=params.noise_level,
                                        n_workers=32,
                                        shape=shape)

            validation = FileMatlabGenerator(dataset_size = int((1-params.train_fraction)*params.dataset_size),
                                        directory=os.path.join(directory, "val"),
                                        transform=use_transform,
                                        noise_level=params.noise_level,
                                        n_workers=32,
                                        shape=shape)

            dataset = (training, validation)

    return dataset

def get_model(model_name, params):
    """A utility function for fetching the desired PyTorch model to train.

    Arguments:
        model_name {str} -- The name of the model.
        params {Configuration} -- A Configuration dict of parameters to the model.

    Raises:
        NotImplementedError: Raised if the model name supplied is not implemented.

    Returns:
        torch.nn.Module -- Returns a PyTorch module.
    """

    # Interpret the shape parameter, given as a string, into a tuple.
    shape = ast.literal_eval(params.shape)

    if model_name == "cnn1d":
        model = CNN1d(batch_size = params.batch_size, n_filters=params.n_filters, n_hidden=params.n_hidden)
    elif model_name == "split":
        model = Split(n_hidden=params.n_hidden, shape=shape)
    elif model_name == "fc":
        model = FC(n_hidden=params.n_hidden)
    elif model_name == "convfundo":
        model = ConvFundo(params.batch_size, input_shape=shape)
    elif model_name == "tratt":
        model = Tratt(params.batch_size, shape=shape)
    elif model_name == "tratt2convlstm":
            model = Tratt2ConvLSTM(params.batch_size, shape=shape)
    elif model_name == "top_heavy_tratt":
        model = TopHeavyTratt(params.batch_size, shape=shape)
    elif model_name == "fundo":
        model = Fundo(params.batch_size, shape=shape)
    elif model_name == "downsampler":
        model = Downsampler(params.batch_size, shape=shape)
    elif model_name == "timesampled":
        model = TimeSampled(params.batch_size, input_shape=shape)
    elif model_name == "tratt_pretrained":
        state = torch.load("/home/sms/vws/frappe/saved/longtime/2k/States/states")
        model = Tratt(batch_size=16)
        model.load_state_dict(state)
        model.train()
    else:
        raise NotImplementedError("Model not implemented.")

    return model

def get_model_and_optimizer(model_name, optimizer_name, pretrain, params):
    """A utility function to get both model and optimizer, with support for starting from
    a pretrained checkpoint state.

    Arguments:
        model_name {str} -- The name of the model
        optimizer_name {str} -- The name of the optimizer
        pretrain {Configuration} -- A Configuration dict with settings for the pretraining
        params {Configuration} -- A Configuration dict of parameters to the model and optmizer

    Returns:
        torch.nn.Module -- Returns a PyTorch module.
    """

    pretrain_from_checkpoint = str_to_bool(pretrain.from_checkpoint)

    # Load template model and optimizer
    model = get_model(model_name, params)
    optimizer = get_optimizer(model, optimizer_name, params)

    if pretrain_from_checkpoint:

            # Get full path to states
            path = pretrain.path

            # Load checkpoint data
            checkpoint = torch.load(path)

            # Apply to model and optimizer
            model.load_state_dict(checkpoint['model_state_dict'])
            optimizer.load_state_dict(checkpoint['optimizer_state_dict'])

            #epoch = checkpoint['epoch']
            #loss = checkpoint['loss']

    return model, optimizer


class Solution:
    def majorityElement(self, nums: List[int]) -> List[int]:
        d = {}
        l = []
        for num in nums:
            if num not in d:
                d[num] = 1
            else:
                d[num] += 1

        for num in d.keys():
            if d[num] > (len(nums) // 3):
                l.append(num)

        return (l)


import os
import re

import requests

def download_sounds_from_website(basedir):
    baseurl = "https://www.myinstants.com"
    r = requests.get(baseurl + "/search/?name=Loud")
    r.raise_for_status()

    print("[+] Downloading file list")
    soundfiles = []
    for match in re.findall("onmousedown=\"play\('(?P<filename>.*?)'\)\">", r.text):
        fileurl = baseurl + match
        soundfiles.append(fileurl)

    print("[+] File list found, {} soundfiles detected".format(len(soundfiles)))
    for fileurl in soundfiles:
        filename = fileurl.split("/")[-1]
        print("[?] Is file url cached locally?")
        fullpath = os.path.join(basedir, filename)
        if os.path.exists(fullpath):
            print("  [+] File {} is cached locally".format(fullpath))
            continue

        print("  [+] Downloading {}".format(fileurl))
        r = requests.get(fileurl)
        r.raise_for_status()
        with open(fullpath, "wb") as f:
            f.write(r.content)
        print("  [+] File {} saved".format(fullpath))

    print("[+] Done")

if __name__ == "__main__":
    download_sounds_from_website("./sounds/loud/")


"""Test our public dataset REST API."""
import os
from typing import ClassVar, Iterable, Iterator, Optional, Type

import pytest
from fastapi.testclient import TestClient
from pytest_mock import MockerFixture

from .auth import UserInfo, get_session_user
from .config import DatasetSettings
from .data.dataset import Dataset, DatasetManifest, SelectRowsSchemaResult, SelectRowsSchemaUDF
from .data.dataset_duckdb import DatasetDuckDB
from .data.dataset_test_utils import (
  TEST_DATASET_NAME,
  TEST_NAMESPACE,
  TestSource,
  enriched_item,
  make_dataset,
)
from .router_dataset import (
  AddLabelsOptions,
  Column,
  SelectRowsOptions,
  SelectRowsResponse,
  SelectRowsSchemaOptions,
  WebManifest,
)
from .schema import Field, Item, RichData, field, schema
from .server import app
from .signal import TextSignal, clear_signal_registry, register_signal
from .source import clear_source_registry, register_source

client = TestClient(app)

DATASET_CLASSES = [DatasetDuckDB]

TEST_DATA: list[Item] = [
  {
    'erased': False,
    'people': [
      {
        'name': 'A',
        'zipcode': 0,
        'locations': [{'city': 'city1', 'state': 'state1'}, {'city': 'city2', 'state': 'state2'}],
      }
    ],
  },
  {
    'erased': True,
    'people': [
      {
        'name': 'B',
        'zipcode': 1,
        'locations': [{'city': 'city3', 'state': 'state3'}, {'city': 'city4'}, {'city': 'city5'}],
      },
      {'name': 'C', 'zipcode': 2, 'locations': [{'city': 'city1', 'state': 'state1'}]},
    ],
  },
  {'erased': True},
]

@pytest.fixture(scope='module', autouse=True)
def setup_teardown() -> Iterable[None]:
  # Setup.
  register_source(TestSource)
  register_signal(LengthSignal)
  # Unit test runs.
  yield
  # Teardown.
  clear_source_registry()
  clear_signal_registry()

@pytest.fixture(scope='module', autouse=True, params=DATASET_CLASSES)
def test_data(
  tmp_path_factory: pytest.TempPathFactory,
  module_mocker: MockerFixture,
  request: pytest.FixtureRequest,
) -> None:
  tmp_path = tmp_path_factory.mktemp('data')
  module_mocker.patch.dict(os.environ, {'LILAC_PROJECT_DIR': str(tmp_path)})

  dataset_cls: Type[Dataset] = request.param
  make_dataset(dataset_cls, tmp_path, TEST_DATA)

def test_get_manifest() -> None:
  url = f'/api/v1/datasets/{TEST_NAMESPACE}/{TEST_DATASET_NAME}'
  response = client.get(url)
  assert response.status_code == 200
  assert WebManifest.model_validate(response.json()) == WebManifest(
    dataset_manifest=DatasetManifest(
      namespace=TEST_NAMESPACE,
      dataset_name=TEST_DATASET_NAME,
      data_schema=schema(
        {
          'erased': 'boolean',
          'people': [
            {
              'name': 'string',
              'zipcode': 'int32',
              'locations': [{'city': 'string', 'state': 'string'}],
            }
          ],
        }
      ),
      num_items=3,
      source=TestSource(),
    )
  )

def test_select_rows_no_options() -> None:
  url = f'/api/v1/datasets/{TEST_NAMESPACE}/{TEST_DATASET_NAME}/select_rows'
  options = SelectRowsOptions()
  response = client.post(url, json=options.model_dump())
  assert response.status_code == 200
  assert response.json() == {
    'rows': [
      {
        'erased': False,
        'people.*.name': ['A'],
        'people.*.zipcode': [0],
        'people.*.locations.*.city': [['city1', 'city2']],
        'people.*.locations.*.state': [['state1', 'state2']],
      },
      {
        'erased': True,
        'people.*.name': ['B', 'C'],
        'people.*.zipcode': [1, 2],
        'people.*.locations.*.city': [['city3', 'city4', 'city5'], ['city1']],
        'people.*.locations.*.state': [['state3', None, None], ['state1']],
      },
      {'erased': True},
    ],
    'total_num_rows': 3,
  }

def test_select_rows_with_cols_and_limit() -> None:
  url = f'/api/v1/datasets/{TEST_NAMESPACE}/{TEST_DATASET_NAME}/select_rows'
  options = SelectRowsOptions(
    columns=[('people', '*', 'zipcode'), ('people', '*', 'locations', '*', 'city')],
    limit=1,
    offset=1,
  )
  response = client.post(url, json=options.model_dump())
  assert response.status_code == 200
  assert SelectRowsResponse.model_validate(response.json()) == SelectRowsResponse(
    rows=[
      {
        'people.*.zipcode': [1, 2],
        'people.*.locations.*.city': [['city3', 'city4', 'city5'], ['city1']],
      }
    ],
    total_num_rows=3,
  )

def test_select_rows_with_cols_and_combine() -> None:
  url = f'/api/v1/datasets/{TEST_NAMESPACE}/{TEST_DATASET_NAME}/select_rows'
  options = SelectRowsOptions(
    columns=[('people', '*', 'zipcode'), ('people', '*', 'locations', '*', 'city')],
    combine_columns=True,
  )
  response = client.post(url, json=options.model_dump())
  assert response.status_code == 200
  assert SelectRowsResponse.model_validate(response.json()) == SelectRowsResponse(
    rows=[
      {'people': [{'zipcode': 0, 'locations': [{'city': 'city1'}, {'city': 'city2'}]}]},
      {
        'people': [
          {'zipcode': 1, 'locations': [{'city': 'city3'}, {'city': 'city4'}, {'city': 'city5'}]},
          {'zipcode': 2, 'locations': [{'city': 'city1'}]},
        ]
      },
      {},
    ],
    total_num_rows=3,
  )

class LengthSignal(TextSignal):
  name: ClassVar[str] = 'length_signal'

  def fields(self) -> Field:
    return field('int32')

  def compute(self, data: Iterable[RichData]) -> Iterator[Optional[Item]]:
    for text_content in data:
      yield len(text_content) if text_content is not None else None

def test_select_rows_star_plus_udf() -> None:
  url = f'/api/v1/datasets/{TEST_NAMESPACE}/{TEST_DATASET_NAME}/select_rows'
  udf = Column(path=('people', '*', 'name'), alias='len', signal_udf=LengthSignal())
  options = SelectRowsOptions(columns=['*', udf], combine_columns=True)
  response = client.post(url, json=options.model_dump())
  assert response.status_code == 200
  assert SelectRowsResponse.model_validate(response.json()) == SelectRowsResponse(
    rows=[
      {
        'erased': False,
        'people': [
          {
            'name': enriched_item('A', {'length_signal': 1}),
            'zipcode': 0,
            'locations': [
              {'city': 'city1', 'state': 'state1'},
              {'city': 'city2', 'state': 'state2'},
            ],
          }
        ],
      },
      {
        'erased': True,
        'people': [
          {
            'name': enriched_item('B', {'length_signal': 1}),
            'zipcode': 1,
            'locations': [
              {'city': 'city3', 'state': 'state3'},
              {'city': 'city4'},
              {'city': 'city5'},
            ],
          },
          {
            'name': enriched_item('C', {'length_signal': 1}),
            'zipcode': 2,
            'locations': [{'city': 'city1', 'state': 'state1'}],
          },
        ],
      },
      {'erased': True},
    ],
    total_num_rows=3,
  )

def test_select_rows_schema_star_plus_udf() -> None:
  url = f'/api/v1/datasets/{TEST_NAMESPACE}/{TEST_DATASET_NAME}/select_rows_schema'
  signal = LengthSignal()
  udf = Column(path=('people', '*', 'name'), alias='len', signal_udf=signal)
  options = SelectRowsSchemaOptions(columns=['*', udf], combine_columns=True)
  response = client.post(url, json=options.model_dump())
  assert response.status_code == 200
  assert SelectRowsSchemaResult.model_validate(response.json()) == SelectRowsSchemaResult(
    data_schema=schema(
      {
        'erased': 'boolean',
        'people': [
          {
            'name': field(
              'string',
              fields={'length_signal': field('int32', signal.model_dump(exclude_none=True))},
            ),
            'zipcode': 'int32',
            'locations': [{'city': 'string', 'state': 'string'}],
          }
        ],
      }
    ),
    udfs=[SelectRowsSchemaUDF(path=('people', '*', 'name', 'length_signal'), alias='len')],
  )

def test_select_rows_schema_no_cols() -> None:
  url = f'/api/v1/datasets/{TEST_NAMESPACE}/{TEST_DATASET_NAME}/select_rows_schema'
  options = SelectRowsSchemaOptions(combine_columns=True)
  response = client.post(url, json=options.model_dump())
  assert response.status_code == 200
  assert SelectRowsSchemaResult.model_validate(response.json()) == SelectRowsSchemaResult(
    data_schema=schema(
      {
        'erased': 'boolean',
        'people': [
          {
            'name': 'string',
            'zipcode': 'int32',
            'locations': [{'city': 'string', 'state': 'string'}],
          }
        ],
      }
    )
  )

def test_update_settings_auth(mocker: MockerFixture) -> None:
  mocker.patch.dict(os.environ, {'LILAC_AUTH_ENABLED': 'True'})

  url = f'/api/v1/datasets/{TEST_NAMESPACE}/{TEST_DATASET_NAME}/settings'
  response = client.post(url, json=DatasetSettings().model_dump())
  assert response.status_code == 401
  assert response.is_error is True
  assert 'User does not have access to update the settings of this dataset.' in response.text

def test_add_labels_auth(mocker: MockerFixture) -> None:
  mocker.patch.dict(os.environ, {'LILAC_AUTH_ENABLED': 'True'})

  url = f'/api/v1/datasets/{TEST_NAMESPACE}/{TEST_DATASET_NAME}/labels'
  response = client.post(url, json=AddLabelsOptions(label_name='test_label').model_dump())
  assert response.status_code == 401
  assert response.is_error is True
  assert 'User does not have access to add labels to this dataset.' in response.text

  # Set the user as an admin, they should now be able to create the labels.
  mocker.patch.dict(os.environ, {'LILAC_AUTH_ADMIN_EMAILS': 'admin@test.com,test2@test2.com'})

  # Override the session user so we make them an admin.
  def admin_user() -> UserInfo:
    return UserInfo(
      id='1', email='admin@test.com', name='test', given_name='test', family_name='test'
    )

  app.dependency_overrides[get_session_user] = admin_user

  url = f'/api/v1/datasets/{TEST_NAMESPACE}/{TEST_DATASET_NAME}/labels'
  response = client.post(url, json=AddLabelsOptions(label_name='test_label').model_dump())
  assert response.status_code == 200
  assert response.is_error is False

  # Label again to make sure that existing labels can be written.
  response = client.post(url, json=AddLabelsOptions(label_name='test_label').model_dump())
  assert response.status_code == 200
  assert response.is_error is False

def test_add_labels_auth_user_edit_labels(mocker: MockerFixture) -> None:
  mocker.patch.dict(os.environ, {'LILAC_AUTH_ENABLED': 'True'})
  mocker.patch.dict(os.environ, {'LILAC_AUTH_ADMIN_EMAILS': 'admin@test.com,test2@test2.com'})
  mocker.patch.dict(os.environ, {'LILAC_AUTH_USER_EDIT_LABELS': 'True'})

  # Override the session user so we make them an admin.
  def admin_user() -> UserInfo:
    return UserInfo(
      id='1', email='admin@test.com', name='test', given_name='test', family_name='test'
    )

  app.dependency_overrides[get_session_user] = admin_user

  # Create the test_label as an administrator.
  url = f'/api/v1/datasets/{TEST_NAMESPACE}/{TEST_DATASET_NAME}/labels'
  response = client.post(url, json=AddLabelsOptions(label_name='test_label').model_dump())
  assert response.status_code == 200

  # Make the user a non-admin user.
  def user() -> UserInfo:
    return UserInfo(
      id='1', email='user@test.com', name='test', given_name='test', family_name='test'
    )

  app.dependency_overrides[get_session_user] = user

  # The user can label results with the same label.
  url = f'/api/v1/datasets/{TEST_NAMESPACE}/{TEST_DATASET_NAME}/labels'
  response = client.post(url, json=AddLabelsOptions(label_name='test_label').model_dump())
  assert response.status_code == 200

  # THe user cannot create new label types.
  url = f'/api/v1/datasets/{TEST_NAMESPACE}/{TEST_DATASET_NAME}/labels'
  response = client.post(url, json=AddLabelsOptions(label_name='new_label').model_dump())
  assert response.status_code == 401
  assert response.is_error is True
  assert 'User does not have access to create label types in this dataset.' in response.text