import json
import os

import requests

from remote_ckan.logs import get_logger

logger = get_logger(__name__)

class RemoteCKAN:
    def __init__(self, url,
                 user_agent='Remote CKAN 1.0',
                 temp_data='checks',
                 requests_timeout=60):
        self.url = url
        self.user_agent = user_agent
        self.errors = []
        self.harvest_sources = {}
        self.organizations = {}
        self.groups = {}
        logger.debug(f'New remote CKAN {url}')
        # save results temp data into a new folder
        self.temp_data = temp_data
        if not os.path.isdir(self.temp_data):
            os.mkdir(self.temp_data)

        self.requests_timeout = requests_timeout

    def set_destination(self, ckan_url, ckan_api_key):
        self.destination_url = ckan_url
        self.api_key = ckan_api_key

    def list_harvest_sources(self,
                             source_type=None,
                             start=0,
                             page_size=100,
                             limit=0,
                             skip_full_source_info=False):
        """ Generator for a list of harvest sources at a CKAN instance
            Params:
                source_type (str): datajson | csw | None=ALL
                limit (int): max number of harvest sources to read
        """

        logger.info(f'List harvest sources {start}-{page_size}')

        package_search_url = f'{self.url}/api/3/action/package_search'
        # TODO use harvest_source_list for harvester ext

        if source_type is None or source_type == 'ALL':
            q = '(type:harvest)'
        else:
            q = f'(type:harvest source_type:{source_type})'

        params = {
            'start': start,
            'rows': page_size,
            'q': q,
            'fq': '+dataset_type:harvest',
            'sort': 'metadata_created asc',
        }
        headers = self.get_request_headers(include_api_key=False)

        logger.debug(f'request {package_search_url} {params}')
        response = requests.get(package_search_url, params=params, headers=headers, timeout=self.requests_timeout)
        if response.status_code >= 400:
            error = f'ERROR getting harvest sources: {response.status_code} {response.text}'
            self.errors.append(error)
            logger.error(error)
            raise ValueError(error)

        data = response.json()

        if not data['success']:
            error = 'ERROR searching harvest sources {}'.format(data['error'])
            logger.error(error)
            self.errors.append(error)
            raise ValueError(error)

        total = data['result']['count']
        count = len(data['result']['results'])
        if count == 0:
            return

        harvest_sources = data['result']['results']
        logger.info(f'{count} ({total}) harvest sources found')

        for hs in harvest_sources:
            total_sources = len(self.harvest_sources)
            if limit > 0 and total_sources >= limit:
                return
            title = hs['title']
            hs_source_type = hs['source_type']  # datajosn, waf, etc
            state = hs['state']
            name = hs['name']
            self.harvest_sources[name] = hs

            logger.info(f'  [{hs_source_type}] Harvest source: {title} [{state}]')
            if state == 'active':
                if not skip_full_source_info:
                    # We don't get full harvest soure info here. We need a custom call
                    source = self.get_full_harvest_source(hs)
                else:
                    logger.info('SKIPPING full havest source info')
                    # faster, but incomplete
                    source = hs

                if source is None:
                    source = {'name': hs['name'], 'created': False, 'updated': False, 'error': True}
                yield source

        # if the page is not full, it is the last one
        if count + 1 < page_size:
            return

        # get next page
        yield from self.list_harvest_sources(
            source_type=source_type,
            start=start + page_size,
            page_size=page_size,
            limit=limit,
            skip_full_source_info=skip_full_source_info)

    def get_full_harvest_source(self, hs, url=None):
        """ get full info (including job status) for a Harvest Source """
        if url is None:
            url = self.url
        harvest_show_url = f'{url}/api/3/action/harvest_source_show'

        if 'id' in hs.keys():
            params = {'id': hs['id']}
        else:
            params = {'name_or_id': hs['name']}

        headers = self.get_request_headers(include_api_key=False)

        error = None
        logger.info(f'Get harvest source data {harvest_show_url} {params}')
        try:
            response = requests.get(harvest_show_url, params=params, headers=headers, timeout=self.requests_timeout)
        except Exception as e:
            error = f'Request harvest source Error: {e}'
        else:
            if response.status_code >= 500:
                error = f'Error [{response.status_code}] getting harvest source info: {harvest_show_url}\n\t{params}'
            elif response.status_code >= 400:
                error = f'Error [{response.status_code}] getting harvest source info: "{response.content}"'

        if error is not None:
            logger.error(error)
            self.errors.append(error)
            name = hs['name']
            source = {'name': name, 'created': False, 'updated': False, 'error': True}
            self.harvest_sources[name] = source
            return None

        full_hs = response.json()
        self.harvest_sources[hs['name']] = full_hs['result']
        self.save_temp_json('harvest-source', full_hs['result']['name'], full_hs['result'])
        return full_hs['result']

    def get_full_package(self, name_or_id, url=None):
        """ get full info (including job status) for a Harvest Source """
        if url is None:
            url = self.url

        show_url = f'{url}/api/3/action/package_show'
        params = {'id': name_or_id}

        headers = self.get_request_headers(include_api_key=True)

        error = None
        logger.info(f'Get package {show_url} {params}')
        try:
            response = requests.get(show_url, params=params, headers=headers, timeout=self.requests_timeout)
        except Exception as e:
            error = f'Request harvest source Error: {e}'
        else:
            if response.status_code >= 500:
                error = f'Error [{response.status_code}] getting harvest source info: {show_url}\n\t{params}'
            elif response.status_code == 404:
                return None
            elif response.status_code >= 400:
                error = f'Error [{response.status_code}] getting harvest source info: "{response.content}"'

        if error is not None:
            logger.error(error)
            self.errors.append(error)
            return None

        package_response = response.json()
        return package_response['result']

    def get_full_organization(self, org, url=None):
        """ get full info (including job status) for a Harvest Source """

        if url is None:
            url = self.url

        organization_show_url = f'{url}/api/3/action/organization_show'

        org_id = org.get('name', org.get('id', None))
        if org_id is None:
            return None

        params = {'id': org_id}
        headers = self.get_request_headers(include_api_key=False)

        logger.info(f'Get organization data {organization_show_url} {params}')

        # not sure why require this
        headers['Content-Type'] = 'application/x-www-form-urlencoded'
        params = json.dumps(params)

        error = None

        try:
            response = requests.post(organization_show_url, data=params, headers=headers, timeout=self.requests_timeout)
        except Exception as e:
            error = f'Request organization Error: {e}'
        else:
            if response.status_code >= 400:
                error = f'Error [{response.status_code}] getting organization info:\n {params} \n{response.headers}\n {response.content}'
                logger.error(error)
                self.errors.append(error)
            else:
                response = response.json()
                org = response['result']

        if error is not None:
            logger.error(error)
            self.errors.append(error)

            return None

        self.save_temp_json('organization', org.get('name', org.get('id', 'unnamed')), org)
        self.organizations[org['name']] = org
        return org

    def get_request_headers(self, include_api_key=True):
        headers = {'User-Agent': self.user_agent}
        if include_api_key:
            headers['X-CKAN-API-Key'] = self.api_key
        return headers

    def create_harvest_source(self, data):
        """ create a harvest source (is just a CKAN dataset/package)
            This will create the organization if it does not already exist.
            params:
                data (dict): Harvest source dict

            returns:
                created (boolean):
                status_code (int): request status code
                error (str): None or error
            """

        org = data.get('organization', None)
        if org is None:
            self.harvest_sources[data['name']].update({'created': False, 'updated': False, 'error': True})
            return False, 0, f'Missing organization at \n\t{data}'

        created, status, error = self.create_organization(data=org)
        if not created:
            self.harvest_sources[data['name']].update({'created': False, 'updated': False, 'error': True})
            return False, status, f'Unable to create organization: {error}'

        config = data.get('config', {})
        if type(config) == str:
            config = json.loads(config)

        default_groups = config.get('default_groups', [])
        if type(default_groups) == str:
            default_groups = [default_groups]
        for group in default_groups:
            created, status, error = self.create_group(group_name=group)
            if not created:
                self.harvest_sources[data['name']].update({'created': False, 'updated': False, 'error': True})
                return False, status, f'Unable to create group: {error}'

        ckan_package = self.get_package_from_data(data)
        self.harvest_sources[data['name']]['ckan_package'] = ckan_package

        package_create_url = f'{self.destination_url}/api/3/action/harvest_source_create'
        logger.info('Creating harvest source {}'.format(ckan_package['title']))

        created, status, error = self.request_ckan(url=package_create_url, method='POST', data=ckan_package)

        if error == 'Already exists':
            # get the ID of the Source (use destination URL)
            del ckan_package['id']
            local_hs = self.get_full_harvest_source(hs=ckan_package, url=self.destination_url)
            ckan_package['id'] = local_hs['id']
            return self.update_harvest_source(data=ckan_package)
        else:
            name = ckan_package['name']
            self.harvest_sources[name].update({'created': created, 'updated': False, 'error': error is not None})
            return created, status, error

    def update_harvest_source(self, data):
        """ update a harvest source if already exists
            params:
                data (dict): Harvest source dict

            returns:
                created (boolean):
                status_code (int): request status code
                error (str): None or error
            """

        package_update_url = f'{self.destination_url}/api/3/action/harvest_source_update'
        logger.info(' ** Updating source {}'.format(data['title']))

        updated, status, error = self.request_ckan(url=package_update_url, method='POST', data=data)
        name = data['name']
        self.harvest_sources[name].update({'created': False, 'updated': updated, 'error': error is not None})

        return updated, status, error

    def get_group_list(self, url=None):
        """ get all groups """
        if url is None:
            url = self.url

        group_list_url = f'{url}/api/3/action/group_list'
        headers = self.get_request_headers(include_api_key=False)
        logger.info(f'Get group list from {group_list_url}')

        error = None

        try:
            response = requests.get(group_list_url, headers=headers, timeout=self.requests_timeout)
        except Exception as e:
            error = f'Request group Error: {e}'
        else:
            if response.status_code >= 400:
                error = f'Error [{response.status_code}] getting group info: \n{response.headers}\n {response.content}'
                logger.error(error)
                self.errors.append(error)
            else:
                response = response.json()
                groups = response['result']

        if error is not None:
            logger.error(error)
            self.errors.append(error)

            return None

        self.save_temp_json('group', 'list', groups)

        return groups

    def get_datasets_in_group(self, group_name):

        logger.info(f'List dataset in group {group_name}')

        package_search_url = f'{self.url}/api/3/action/package_search'

        params = {
            'start': 0,
            'rows': 500,
            'fq': 'groups:{}'.format(group_name)
        }

        headers = self.get_request_headers(include_api_key=True)

        logger.debug(f'request {package_search_url} {params}')
        response = requests.get(package_search_url, params=params, headers=headers, timeout=self.requests_timeout)
        if response.status_code >= 400:
            error = f'ERROR getting dataset in groups: {response.status_code} {response.text}'
            self.errors.append(error)
            logger.error(error)
            raise ValueError(error)

        data = response.json()

        if not data['success']:
            error = 'ERROR searching datasets in group {}'.format(data['error'])
            logger.error(error)
            self.errors.append(error)
            raise ValueError(error)

        return data['result']['results']

    def get_full_group(self, group_name, url=None, include_datasets=False):
        """ get full info (including job status) for a Harvest Source """
        if url is None:
            url = self.url

        group_show_url = f'{url}/api/3/action/group_show'

        params = {'id': group_name}
        if include_datasets:
            params['include_datasets'] = True

        headers = self.get_request_headers(include_api_key=False)

        logger.info(f'Get group data {group_show_url} {params}')

        # not sure why require this
        headers['Content-Type'] = 'application/x-www-form-urlencoded'
        params = json.dumps(params)

        error = None

        try:
            response = requests.post(group_show_url, data=params, headers=headers, timeout=self.requests_timeout)
        except Exception as e:
            error = f'Request group Error: {e}'
        else:
            if response.status_code >= 400:
                error = f'Error [{response.status_code}] getting group info:\n {params} \n{response.headers}\n {response.content}'
                logger.error(error)
                self.errors.append(error)
            else:
                response = response.json()
                group = response['result']

        if error is not None:
            logger.error(error)
            self.errors.append(error)

            return None

        self.save_temp_json('group', group.get('name', group.get('id', 'unnamed')), group)
        self.groups[group['name']] = group
        return group

    def create_group(self, group_name):
        """ Creates a new group in CKAN destination """

        full_group = self.get_full_group(group_name=group_name)
        if full_group is None:
            return None
        # remove packages from data because they are not harvested yet
        full_group.pop('packages', None)
        full_group.pop('id', None)

        logger.info('Creating group {}'.format(full_group['name']))

        group_create_url = f'{self.destination_url}/api/3/action/group_create'

        # check if group exists at destination url

        group_at_destination = self.get_full_group(group_name=group_name, url=self.destination_url)

        if group_at_destination is not None:
            full_group['id'] = group_at_destination['id']
            created, status, error = self.update_group(data=full_group)
        else:
            created, status, error = self.request_ckan(method='POST', url=group_create_url, data=full_group)

        return created, status, error

    def update_group(self, data):
        """ update an existing group """

        group_update_url = f'{self.destination_url}/api/3/action/group_update'
        return self.request_ckan(method='POST', url=group_update_url, data=data)

    def create_organization(self, data):
        """ Creates a new organization in CKAN destination
            Params:
                data (dics): Required fields to create
        """

        # get more info about org, we miss organization extras
        # We need the organization_type in some cases
        full_org = self.get_full_organization(org=data)
        if full_org is None:
            return False, 0, f'Failed to get organization for {data}'
        extras = data.get('extras', [])
        for extra in full_org.get('extras', []):
            if extra.get('state', 'active') == 'active':
                extras.append({'key': extra['key'], 'value': extra['value']})
                logger.info('New extra found at org {}={}'.format(extra['key'], extra['value']))

        org_create_url = f'{self.destination_url}/api/3/action/organization_create'

        organization = {
            'name': data['name'],
            'title': data['title'],
            'description': data['description'],
            'image_url': data['image_url'],
            'extras': extras
        }
        logger.info('Creating organization {}'.format(data['name']))

        organization_at_destination = self.get_full_organization(org=data, url=self.destination_url)

        if organization_at_destination is not None:
            organization['id'] = organization_at_destination['id']
            created, status, error = self.update_organization(data=organization)
        else:
            created, status, error = self.request_ckan(method='POST', url=org_create_url, data=organization)

        return created, status, error

    def update_organization(self, data):
        """ update an existing organization """

        org_update_url = f'{self.destination_url}/api/3/action/organization_update'
        return self.request_ckan(method='POST', url=org_update_url, data=data)

    def get_config(self, data):
        """ get config and extras from full data package and return a final str config """
        config = data.get('config', {})
        if type(config) == str:
            config = json.loads(config)

        # We may have config defined as extras
        extras = data.get('extras', [])

        for extra in extras:
            if extra['key'] == 'config':
                logger.info(f'Config found in extras: {extra}')
                value = json.loads(extra['value'])
                config.update(value)
                extras.remove(extra)

        # CKAN drop config values as boolean, move to string
        for k, v in config.items():
            if type(v) is bool:
                config[k] = str(v)
            # DOI CKAN harvest source is setting force_all as a string and CKAN expects a boolean
            if k == "force_all":
                if v.lower() == "true":
                    config[k] = True
                elif v.lower() == "false":
                    config[k] = False
        return json.dumps(config)

    def request_ckan(self, method, url, data):
        """ request CKAN and get results """

        headers = self.get_request_headers(include_api_key=True)
        # do not modify original data
        data2 = data.copy()

        try:
            if method == 'POST':
                req = requests.post(url, json=data2, headers=headers, timeout=self.requests_timeout)
            elif method == 'GET':
                req = requests.get(url, params=data2, headers=headers, timeout=self.requests_timeout)
            else:
                raise ValueError(f'Invalid method {method}')

        except Exception as e:
            error = 'ERROR at {} {}: {}'.format(url, data2, e)
            return False, 0, error

        content = req.content
        try:
            json_content = json.loads(content)
        except Exception as e:
            error = 'ERROR parsing JSON: {} {}: {}'.format(url, content, e)
            logger.error(error)
            self.errors.append(error)
            return False, 0, error

        if req.status_code >= 400:
            if 'already exists' in str(content) or 'already in use' in str(content):
                return False, req.status_code, 'Already exists'
            error = ('ERROR status: {}'
                     '\n\t content:{}'
                     '\n\t sent: {}'.format(req.status_code, req.content, data2))
            self.errors.append(error)
            logger.error(error)
            return False, req.status_code, error

        if not json_content['success']:
            error = 'API response failed: {}'.format(json_content.get('error', None))
            logger.error(error)
            self.errors.append(error)
            return False, req.status_code, error

        logger.info('Request OK {}'.format(url))
        return True, req.status_code, None

    def get_package_from_data(self, data):
        """ get full data package and return a final CKAN package """

        config = self.get_config(data)
        extras = data.get('extras', [])

        ret = {
            'name': data['name'],
            'owner_org': data['organization']['name'],
            'title': data['title'],
            'url': data['url'],
            'notes': data['notes'],
            'source_type': data['source_type'],
            'frequency': data['frequency'],
            'config': config,
            'extras': extras
        }

        if data.get('database', None) is not None:
            ret['database'] = data['database']

        if data.get('port', None) is not None:
            ret['port'] = data['port']

        if data.get('extra_search_criteria', None) is not None:
            ret['extra_search_criteria'] = data['extra_search_criteria']

        cfg = json.loads(config)

        if cfg.get('collection_metadata_url', None) is not None:
            ret['collection_metadata_url'] = cfg['collection_metadata_url']

        if cfg.get('database', None) is not None:
            ret['database'] = cfg['database']

        if cfg.get('port', None) is not None:
            ret['port'] = cfg['port']

        if cfg.get('extra_search_criteria', None) is not None:
            ret['extra_search_criteria'] = cfg['extra_search_criteria']

        if data.get('id', None) is not None:
            ret['id'] = data['id']

        if cfg.get('validator_profiles', None) is not None:
            ret['validator_profiles'] = cfg['validator_profiles']
            # remove this from here, Spatial ext will try to
            # validate and we are using custom profiles
            cfg.pop('validator_profiles')
            ret['config'] = json.dumps(cfg)

        if cfg.get('private_datasets', None) is not None:
            ret['private_datasets'] = cfg['private_datasets']

        return ret

    def save_temp_json(self, data_type, data_name, data):
        filename = '{}-{}.json'.format(data_type, data_name)
        path = os.path.join(self.temp_data, filename)
        f = open(path, 'w')
        f.write(json.dumps(data, indent=4))
        f.close()
        logger.info('{} {} saved at /{}'.format(data_type, data_name, self.temp_data))


def process_file(file_name):
    """
    Given a file name, returns a list of lists [name, gender, births]

    :param file_name: a string
    :return: a list of lists, [[name1, gender1, births1], [name2, gender2, births2]...]

    Example:
    process_file('babynames/yob1880.txt')
        will return
    [["Mary","F",7065], ["Anna","F",2604],...]

    """
    list_of_name_details = []
    with open(file_name) as f:
        for line in f:
            list_of_name_details.append(line.strip().split(','))

    return list_of_name_details

def total_births(year):
    """
    :param year: an integer, between 1880 and 2010
    :return: an integer, the total births of all the babies in that year
    """
    proper_file = 'babynames/yob'+str(year)+'.txt'
    proper_details = process_file(proper_file)
    births = 0

    for i in proper_details:
        births += int(i[2])

    return births

def proportion(name, gender, year):
    """

    :param name: a string, first name
    :param gender: a string, "F" or "M"
    :param year: an integer, between 1880 and 2010
    :return: a float number, the proportion of babies with the given name and given gender to total births in given year
    """
    proper_file = 'babynames/yob'+str(year)+'.txt'
    proper_details = process_file(proper_file)
    total_babies = total_births(year)
    babies_with_name = 0

    for i in proper_details:
        if i[0] == name.capitalize() and i[1] == gender:
            babies_with_name = i[2]

    return int(babies_with_name)/total_babies

def highest_year(name, gender):
    """

    :param name: a string
    :param gender: a string, "F" or "M"
    :return: an integer, the year when the given name has the highest proportion over the years (among all the
    proportions of the same name with same gender in all years)
    """
    years = range(1880, 2011)
    year_with_highest_proportion = 0

    for year in years:
        year_proportion = proportion(name, gender, year)
        if year_proportion > year_with_highest_proportion:
            highest_year = year
            year_with_highest_proportion = year_proportion

    return highest_year

def main():
    print("Let's find out which year my name had the highest proportion among all names with the same gender")
    my_name = input('My name is: ')
    my_gender = input("My gender is (M or F): ")

    print(f'The year in which my name, {my_name}, had the highest proportion among all names with the same gender, {my_gender}, in all years from 1880 to 2010:')
    print(highest_year(my_name, my_gender))
    input("Press enter to close program")

if __name__ == '__main__':
    main()


#!/usr/bin/python
# -*- coding: utf-8 -*-
"""Update Schema from 2.0 to 2.1."""
from playhouse.migrate import SchemaMigrator, migrate
from ..globals import DB

def update_schema():
    """Update schema from 2.0 to 2.1."""
    DB.execute_sql('alter table projects rename constraint proposals_pkey to projects_pkey')
    rename_indexes = {
        'atoolproject_proposal_id': 'atoolproject_project_id',
        'citationproject_proposal_id': 'citationproject_project_id',
        'projectgroup_proposal_id': 'projectgroup_project_id',
        'projectinstrument_proposal_id': 'projectinstrument_project_id',
        'projectparticipant_proposal_id': 'projectparticipant_project_id',
        'proposals_accepted_date': 'projects_accepted_date',
        'proposals_actual_end_date': 'projects_actual_end_date',
        'proposals_actual_start_date': 'projects_actual_start_date',
        'proposals_closed_date': 'projects_closed_date',
        'proposals_created': 'projects_created',
        'proposals_deleted': 'projects_deleted',
        'proposals_short_name': 'projects_short_name',
        'proposals_submitted_date': 'projects_submitted_date',
        'proposals_suspense_date': 'projects_suspense_date',
        'proposals_title': 'projects_title',
        'proposals_updated': 'projects_updated',
        'transsap_proposal_id': 'transsap_project_id',
        'transsip_proposal_id': 'transsip_project_id'
    }
    for old_index, new_index in rename_indexes.items():
        DB.execute_sql('alter index {} rename to {}'.format(old_index, new_index))
    migrator = SchemaMigrator(DB)
    for table_name in ['transsip', 'transsap']:
        for col_name in ['created', 'deleted', 'updated']:
            migrate(migrator.add_index(table_name, (col_name,), False))


import json
import datetime
from io import BytesIO
from urllib import request

import requests
from django.utils.datastructures import MultiValueDictKeyError
from drf_yasg import openapi
from drf_yasg.utils import swagger_auto_schema
from rest_framework import viewsets, permissions, status
from rest_framework.generics import get_object_or_404
from rest_framework.parsers import JSONParser, MultiPartParser
from rest_framework.response import Response
from rest_framework.views import APIView
from fokus.models import Analysis, AnalysisImage
from fokus.serializers import AnalysisSerializer, AnalysisImageSerializer, PUTImageToSessionRequestSerializer, \
    PUTImageToSessionResponseSerializer, POSTCreateNewSessionRequestSerializer, POSTCreateNewSessionResponseSerializer, \
    StopSessionRequestSerializer, StopSessionResponseSerializer

# Create your views here.
class AnalysisList(APIView):
    @swagger_auto_schema(
        responses={200: AnalysisSerializer(many=True)}
    )
    def get(self, _request, format=None):
        analyses = Analysis.objects.filter(ongoing=False, user=_request.user)
        serializer = AnalysisSerializer(analyses, many=True)
        return Response(serializer.data)

    @swagger_auto_schema(
        request_body=POSTCreateNewSessionRequestSerializer,
        responses={201: POSTCreateNewSessionResponseSerializer(many=False)}
    )
    def post(self, _request, format=None):
        data = _request.data
        if data['start'] == 'true':

            analysis = None

            try:
                analysis = Analysis.objects.get(ongoing=True, user=_request.user)
            except Analysis.DoesNotExist:
                analysis = Analysis(user=_request.user,
                                    time_started=datetime.datetime.now(),
                                    description=data['description'],
                                    ongoing=True,
                                    )

                Analysis.save(analysis)
            except Analysis.MultipleObjectsReturned:
                for analysis in Analysis.objects.filter(ongoing=True, user=_request.user):
                    analysis.end_session()
            except Exception as e:
                return Response({"message": str(e)}, status=status.HTTP_500_INTERNAL_SERVER_ERROR)

            response = {
                "error": "false",
                "message": "New session successfully created!",
                "description": analysis.description,
                "id": analysis.id,
                "ongoing": analysis.ongoing,
            }

            return Response(response)

        else:
            return Response({"message": "Invalid request"}, status=status.HTTP_400_BAD_REQUEST)

    @swagger_auto_schema(
        request_body=StopSessionRequestSerializer,
        responses={200: StopSessionResponseSerializer(many=False)},
    )
    def put(self, _request):
        try:
            analysis = Analysis.objects.get(user=_request.user, ongoing=True)
            return AnalysisDetail.as_view()(_request._request, pk=analysis.id)
        except Analysis.DoesNotExist:
            return Response({'message': 'No currently ongoing session'}, status=status.HTTP_400_BAD_REQUEST)
        except Analysis.MultipleObjectsReturned:
            analysis_list = Analysis.objects.filter(ongoing=True, user=_request.user).order_by('time_started')
            print(analysis_list)
            for analysis in analysis_list[:len(analysis_list) - 1]:
                analysis.end_session()
                analysis.save()
            return self.put(_request)

class AnalysisDetail(APIView):
    parser_classes = [MultiPartParser, JSONParser]
    @swagger_auto_schema(
        responses={200: AnalysisSerializer(many=False)},
    )
    def get(self, _request, pk, format=None):
        analysis = get_object_or_404(Analysis.objects.filter(user=_request.user), pk=pk)
        serializer = AnalysisSerializer(analysis)

        return Response(serializer.data)

    @swagger_auto_schema(
        request_body=PUTImageToSessionRequestSerializer,
        responses={200: PUTImageToSessionResponseSerializer(many=False)},
    )
    def put(self, _request, pk):
        data = _request.data
        analysis = get_object_or_404(Analysis.objects.filter(user=_request.user), pk=pk)
        if data['ongoing'] == 'false':
            # TODO: handle if session is already stopped
            analysis.end_session()
            Analysis.save(analysis)
            response = {
                "message": "Session stopped and saved!",
                "id": analysis.id,
                "ongoing": analysis.ongoing,
                "time_started": analysis.time_started,
                "session_length": str(analysis.session_length).split(".")[0],
                "focus_length": str(analysis.focus_length).split(".")[0]
            }

            return Response(response)
        else:
            try:
                # TODO: handle if session is already stopped
                _file = _request.data['file']
                image = AnalysisImage(analysis_session=analysis, image=_file)
                AnalysisImage.save(image)

                predict_response = requests.post('https://fokusin-model-ejh5i5qlpq-et.a.run.app/predict',
                                                 json={'image_url': image.image.url})

                confidence = predict_response.json()['predict_image']

                image.status = True if confidence > 0.50 else False
                image.confidence = confidence * 100
                image.save()

                response = {
                    "analysis_id": analysis.id,
                    "focus": image.status,
                    "confidence": image.confidence
                }

                return Response(response)
            except MultiValueDictKeyError:
                return Response({"message": "No image supplied"}, status=status.HTTP_400_BAD_REQUEST)
            except Exception as e:
                return Response({"message": str(e)}, status=status.HTTP_500_INTERNAL_SERVER_ERROR)


from brownie import accounts, UnsafeBank, Attacker

account = accounts[0]

def main():

    bank = Bank.deploy({"from": account})

    attacker = Attacker.deploy(bank.address, {"from": account})

    amount_to_deposit = 1000000000000000000
    amount_from_attacker = amount_to_deposit / 2

    print(f"Bank initial balance {bank.balance()}")

    # Deposit some funds in bank
    bank.deposit({"from": account, "value": amount_to_deposit})

    print(f"Bank balance after deposit {bank.balance()}")

    # Call start function from attacker
    attacker.start({"from": account, "value": amount_from_attacker})

    print(f"Bank balance after attack {bank.balance()}")

    # Assert bank funds go to zero
    assert(bank.balance() == 0)
    assert(attacker.balance() == amount_to_deposit + amount_from_attacker)


# -*- coding: utf-8 -*-
# File: base.py

import os
import inspect
import pprint
from collections import namedtuple
import weakref

from ...utils.argtools import log_once
from ...utils.utils import get_rng
from ...utils.develop import deprecated
from ..image import check_dtype

# Cannot import here if we want to keep backward compatibility.
# Because this causes circular dependency
# from .transform import TransformList, PhotometricTransform, TransformFactory

__all__ = ['Augmentor', 'ImageAugmentor', 'AugmentorList', 'PhotometricAugmentor']

def _reset_augmentor_after_fork(aug_ref):
    aug = aug_ref()
    if aug:
        aug.reset_state()

def _default_repr(self):
    """
    Produce something like:
    "imgaug.MyAugmentor(field1={self.field1}, field2={self.field2})"

    It assumes that the instance `self` contains attributes that match its constructor.
    """
    classname = type(self).__name__
    argspec = inspect.getfullargspec(self.__init__)
    assert argspec.varargs is None, "The default __repr__ in {} doesn't work for varargs!".format(classname)
    assert argspec.varkw is None, "The default __repr__ in {} doesn't work for kwargs!".format(classname)
    defaults = {}

    fields = argspec.args[1:]
    defaults_pos = argspec.defaults
    if defaults_pos is not None:
        for f, d in zip(fields[::-1], defaults_pos[::-1]):
            defaults[f] = d

    for k in argspec.kwonlyargs:
        fields.append(k)
        if k in argspec.kwonlydefaults:
            defaults[k] = argspec.kwonlydefaults[k]

    argstr = []
    for f in fields:
        assert hasattr(self, f), \
            "Attribute {} in {} not found! Default __repr__ only works if " \
            "the instance has attributes that match the constructor.".format(f, classname)
        attr = getattr(self, f)
        if f in defaults and attr is defaults[f]:
            continue
        argstr.append("{}={}".format(f, pprint.pformat(attr)))
    return "imgaug.{}({})".format(classname, ', '.join(argstr))

ImagePlaceholder = namedtuple("ImagePlaceholder", ["shape"])

class ImageAugmentor(object):
    """
    Base class for an augmentor

    ImageAugmentor should take images of type uint8 in range [0, 255], or
    floating point images in range [0, 1] or [0, 255].

    Attributes:
        rng: a numpy :class:`RandomState`
    """

    def __init__(self):
        self.reset_state()

        # only available on Unix after Python 3.7
        if hasattr(os, 'register_at_fork'):
            os.register_at_fork(
                after_in_child=lambda: _reset_augmentor_after_fork(weakref.ref(self)))

    def _init(self, params=None):
        if params:
            for k, v in params.items():
                if k != 'self' and not k.startswith('_'):
                    setattr(self, k, v)

    def reset_state(self):
        """
        Reset rng and other state of the augmentor.

        Similar to :meth:`DataFlow.reset_state`, the caller of Augmentor
        is responsible for calling this method (once or more times) in the **process that uses the augmentor**
        before using it.

        If you use a built-in augmentation dataflow (:class:`AugmentImageComponent`, etc),
        this method will be called in the dataflow's own `reset_state` method.

        If you use Python≥3.7 on Unix, this method will be automatically called after fork,
        and you do not need to bother calling it.
        """
        self.rng = get_rng(self)

    def _rand_range(self, low=1.0, high=None, size=None):
        """
        Generate uniform float random number between low and high using `self.rng`.
        """
        if high is None:
            low, high = 0, low
        if size is None:
            size = []
        return self.rng.uniform(low, high, size).astype("float32")

    def __str__(self):
        try:
            return _default_repr(self)
        except AssertionError as e:
            log_once(e.args[0], 'warn')
            return super(Augmentor, self).__repr__()

    __repr__ = __str__

    def get_transform(self, img):
        """
        Instantiate a :class:`Transform` object to be used given the input image.
        Subclasses should implement this method.

        The :class:`ImageAugmentor` often has random policies which generate deterministic transform.
        Any of those random policies should happen inside this method and instantiate
        an actual deterministic transform to be performed.
        The returned :class:`Transform` object should perform deterministic transforms
        through its :meth:`apply_*` method.

        In this way, the returned :class:`Transform` object can be used to transform not only the
        input image, but other images or coordinates associated with the image.

        Args:
            img (ndarray): see notes of this class on the requirements.

        Returns:
            Transform
        """
        # This should be an abstract method
        # But we provide an implementation that uses the old interface,
        # for backward compatibility
        log_once("The old augmentor interface was deprecated. "
                 "Please implement {} with `get_transform` instead!".format(self.__class__.__name__),
                 "warning")

        def legacy_augment_coords(self, coords, p):
            try:
                return self._augment_coords(coords, p)
            except AttributeError:
                pass
            try:
                return self.augment_coords(coords, p)
            except AttributeError:
                pass
            return coords  # this is the old default

        p = None  # the default return value for this method
        try:
            p = self._get_augment_params(img)
        except AttributeError:
            pass
        try:
            p = self.get_augment_params(img)
        except AttributeError:
            pass

        from .transform import BaseTransform, TransformFactory
        if isinstance(p, BaseTransform):  # some old augs return Transform already
            return p

        return TransformFactory(name="LegacyConversion -- " + str(self),
                                apply_image=lambda img: self._augment(img, p),
                                apply_coords=lambda coords: legacy_augment_coords(self, coords, p))

    def augment(self, img):
        """
        Create a transform, and apply it to augment the input image.

        This can save you one line of code, when you only care the augmentation of "one image".
        It will not return the :class:`Transform` object to you
        so you won't be able to apply the same transformation on
        other data associated with the image.

        Args:
            img (ndarray): see notes of this class on the requirements.

        Returns:
            img: augmented image.
        """
        check_dtype(img)
        t = self.get_transform(img)
        return t.apply_image(img)

    # ###########################
    # Legacy interfaces:
    # ###########################
    @deprecated("Please use `get_transform` instead!", "2020-06-06", max_num_warnings=3)
    def augment_return_params(self, d):
        t = self.get_transform(d)
        return t.apply_image(d), t

    @deprecated("Please use `transform.apply_image` instead!", "2020-06-06", max_num_warnings=3)
    def augment_with_params(self, d, param):
        return param.apply_image(d)

    @deprecated("Please use `transform.apply_coords` instead!", "2020-06-06", max_num_warnings=3)
    def augment_coords(self, coords, param):
        return param.apply_coords(coords)

class AugmentorList(ImageAugmentor):
    """
    Augment an image by a list of augmentors
    """

    def __init__(self, augmentors):
        """
        Args:
            augmentors (list): list of :class:`ImageAugmentor` instance to be applied.
        """
        assert isinstance(augmentors, (list, tuple)), augmentors
        self.augmentors = augmentors
        super(AugmentorList, self).__init__()

    def reset_state(self):
        """ Will reset state of each augmentor """
        super(AugmentorList, self).reset_state()
        for a in self.augmentors:
            a.reset_state()

    def get_transform(self, img):
        check_dtype(img)
        assert img.ndim in [2, 3], img.ndim

        from .transform import LazyTransform, TransformList
        # The next augmentor requires the previous one to finish.
        # So we have to use LazyTransform
        tfms = []
        for idx, a in enumerate(self.augmentors):
            if idx == 0:
                t = a.get_transform(img)
            else:
                t = LazyTransform(a.get_transform)

            if isinstance(t, TransformList):
                tfms.extend(t.tfms)
            else:
                tfms.append(t)
        return TransformList(tfms)

    def __str__(self):
        repr_each_aug = ",\n".join(["  " + repr(x) for x in self.augmentors])
        return "imgaug.AugmentorList([\n{}])".format(repr_each_aug)

    __repr__ = __str__

Augmentor = ImageAugmentor
"""
Legacy name. Augmentor and ImageAugmentor are now the same thing.
"""

class PhotometricAugmentor(ImageAugmentor):
    """
    A base class for ImageAugmentor which only affects pixels.

    Subclass should implement `_get_params(img)` and `_impl(img, params)`.
    """
    def get_transform(self, img):
        p = self._get_augment_params(img)
        from .transform import PhotometricTransform
        return PhotometricTransform(func=lambda img: self._augment(img, p),
                                    name="from " + str(self))

    def _get_augment_params(self, _):
        return None


import os
from PySide6.QtCore import Slot, QItemSelection, Qt
from PySide6.QtGui import QAction
from spinetoolbox.helpers import open_url
from spinetoolbox.mvcmodels.file_list_models import FileListModel
from spine_engine.config import TOOL_OUTPUT_DIR
from spine_engine.project_item.project_item_resource import CmdLineArg, make_cmd_line_arg, LabelArg
from spine_engine.utils.helpers import resolve_python_interpreter
from .commands import (
    UpdateKillCompletedProcesses,
    UpdateToolExecuteInWorkCommand,
    UpdateToolOptionsCommand,
    UpdateLogProcessOutput,
)
from ..db_writer_item_base import DBWriterItemBase
from ..commands import UpdateCmdLineArgsCommand, UpdateGroupIdCommand
from .item_info import ItemInfo
from .widgets.custom_menus import ToolSpecificationMenu
from .widgets.options_widgets import JuliaOptionsWidget
from .executable_item import ExecutableItem
from .utils import flatten_file_path_duplicates, find_file
from ..models import ToolCommandLineArgsModel
from .output_resources import scan_for_resources

class Tool(DBWriterItemBase):
    def __init__(
        self,
        name,
        description,
        x,
        y,
        toolbox,
        project,
        specification_name="",
        execute_in_work=True,
        cmd_line_args=None,
        options=None,
        kill_completed_processes=False,
        log_process_output=False,
        group_id=None,
    ):
        """Tool class.

        Args:
            name (str): Object name
            description (str): Object description
            x (float): Initial X coordinate of item icon
            y (float): Initial Y coordinate of item icon
            toolbox (ToolboxUI): QMainWindow instance
            project (SpineToolboxProject): the project this item belongs to
            specification_name (str): Name of this Tool's Tool specification
            execute_in_work (bool): Execute associated Tool specification in work (True) or source directory (False)
            cmd_line_args (list, optional): Tool command line arguments
            options (dict, optional): misc tool options. At the moment it just holds the location of the julia sysimage
            kill_completed_processes (bool): whether to kill completed persistent processes
            log_process_output (bool): whether to log process output to a file
            group_id (str, optional): execution group id
        """
        super().__init__(name, description, x, y, project)
        self._toolbox = toolbox
        self.execute_in_work = execute_in_work
        if cmd_line_args is None:
            cmd_line_args = []
        self.cmd_line_args = cmd_line_args
        self._cmdline_args_model = ToolCommandLineArgsModel(self)
        self._specification = self._project.get_specification(specification_name)
        if specification_name and not self._specification:
            self._logger.msg_error.emit(
                f"Tool {self.name} should have a Tool specification "
                f"{specification_name} but it was not found"
            )
        self._group_id = group_id
        self._cmdline_args_model.args_updated.connect(self._push_update_cmd_line_args_command)
        self._populate_cmdline_args_model()
        self._input_file_model = FileListModel(header_label="Available resources", draggable=True)
        # Make directory for results
        self.output_dir = os.path.join(self.data_dir, TOOL_OUTPUT_DIR)
        self._specification_menu = None
        self._options = options if options is not None else {}
        self._kill_completed_processes = kill_completed_processes
        self._log_process_output = log_process_output
        self._resources_from_upstream = list()
        self._resources_from_downstream = list()
        self._available_resources = list()
        self._req_file_paths = set()
        self._input_files_not_found = None

    def set_up(self):
        execute_in_work = self.execute_in_work
        super().set_up()
        self.do_update_execution_mode(execute_in_work)

    @property
    def group_id(self):
        return self._group_id

    def _get_options_widget(self):
        """Returns a widget to specify the options for this tool.
        It is embedded in the ui in ``self._update_tool_ui()``.

        Returns:
            OptionsWidget
        """
        # At the moment only Julia has options, but the code is made generic
        constructors = {"julia": JuliaOptionsWidget}  # Add others as needed
        tooltype = self.specification().tooltype
        constructor = constructors.get(tooltype)
        if constructor is None:
            return None
        if tooltype not in self._properties_ui.options_widgets:
            self._properties_ui.options_widgets[tooltype] = constructor()
        options_widget = self._properties_ui.options_widgets[tooltype]
        options_widget.set_tool(self)
        options_widget.do_update_options(self._options)
        return options_widget

    @staticmethod
    def item_type():
        """See base class."""
        return ItemInfo.item_type()

    @staticmethod
    def item_category():
        """See base class."""
        return ItemInfo.item_category()

    def make_signal_handler_dict(self):
        """Returns a dictionary of all shared signals and their handlers.
        This is to enable simpler connecting and disconnecting."""
        s = super().make_signal_handler_dict()
        s[self._properties_ui.toolButton_tool_specification.clicked] = self.show_specification_window
        s[self._properties_ui.pushButton_tool_results.clicked] = self._open_results_directory
        s[self._properties_ui.comboBox_tool.textActivated] = self.update_specification
        s[self._properties_ui.radioButton_execute_in_work.toggled] = self.update_execution_mode
        s[self._properties_ui.toolButton_add_file_path_arg.clicked] = self._add_selected_file_path_args
        s[self._properties_ui.toolButton_remove_arg.clicked] = self._remove_arg
        s[
            self._properties_ui.treeView_input_files.selectionModel().selectionChanged
        ] = self._update_add_args_button_enabled
        s[
            self._properties_ui.treeView_cmdline_args.selectionModel().selectionChanged
        ] = self._update_remove_args_button_enabled
        s[self._properties_ui.lineEdit_group_id.editingFinished] = self._set_group_id
        s[self._properties_ui.kill_consoles_check_box.clicked] = self._set_kill_completed_processes
        s[self._properties_ui.log_process_output_check_box.clicked] = self._set_log_process_output
        return s

    @Slot()
    def _set_group_id(self):
        """Pushes a command to update group id whenever the user edits the line edit."""
        group_id = self._properties_ui.lineEdit_group_id.text()
        if not group_id:
            group_id = None
        if self._group_id == group_id:
            return
        self._toolbox.undo_stack.push(UpdateGroupIdCommand(self, group_id))

    def do_set_group_id(self, group_id):
        """Sets group id."""
        self._group_id = group_id
        if self._active:
            self._properties_ui.lineEdit_group_id.setText(group_id)

    @Slot(QItemSelection, QItemSelection)
    def _update_add_args_button_enabled(self, _selected, _deselected):
        self._do_update_add_args_button_enabled()

    def _do_update_add_args_button_enabled(self):
        enabled = self._properties_ui.treeView_input_files.selectionModel().hasSelection()
        self._properties_ui.toolButton_add_file_path_arg.setEnabled(enabled)

    @Slot(QItemSelection, QItemSelection)
    def _update_remove_args_button_enabled(self, _selected, _deselected):
        self._do_update_remove_args_button_enabled()

    def _do_update_remove_args_button_enabled(self):
        enabled = self._properties_ui.treeView_cmdline_args.selectionModel().hasSelection()
        self._properties_ui.toolButton_remove_arg.setEnabled(enabled)

    def _update_kill_consoles_check_box_enabled(self):
        """Updates the enabled state of 'Kill consoles after execution' check box."""
        enabled = self._specification is None or self._specification.tooltype in ("python", "julia")
        self._properties_ui.kill_consoles_check_box.setEnabled(enabled)

    def restore_selections(self):
        """Restore selections into shared widgets when this project item is selected."""
        self._properties_ui.treeView_input_files.setModel(self._input_file_model)
        self._properties_ui.treeView_cmdline_args.setModel(self._cmdline_args_model)
        self._properties_ui.treeView_cmdline_args.expandAll()
        self.update_execute_in_work_button()
        self._properties_ui.label_jupyter.elided_mode = Qt.ElideMiddle
        self._properties_ui.label_jupyter.hide()
        self._update_tool_ui()
        self._do_update_add_args_button_enabled()
        self._do_update_remove_args_button_enabled()
        self._properties_ui.lineEdit_group_id.setText(self._group_id)

    @Slot(bool)
    def show_specification_window(self, _=True):
        """Opens the settings window."""
        self._toolbox.show_specification_form(self.item_type(), self.specification(), self)

    @Slot(bool)
    def update_execution_mode(self, checked):
        """Pushes a new UpdateToolExecuteInWorkCommand to the toolbox stack."""
        self._toolbox.undo_stack.push(UpdateToolExecuteInWorkCommand(self, checked))

    def do_update_execution_mode(self, execute_in_work):
        """Updates execute_in_work setting."""
        if self.execute_in_work == execute_in_work:
            return
        self.execute_in_work = execute_in_work
        self.update_execute_in_work_button()

    def update_execute_in_work_button(self):
        """Sets execute in work radio button check state according to
        execute_in_work instance variable."""
        if not self._active:
            return
        self._properties_ui.radioButton_execute_in_work.blockSignals(True)
        self._properties_ui.radioButton_execute_in_source.blockSignals(True)
        if self.execute_in_work:
            self._properties_ui.radioButton_execute_in_work.setChecked(True)
        else:
            self._properties_ui.radioButton_execute_in_source.setChecked(True)
        self._properties_ui.radioButton_execute_in_work.blockSignals(False)
        self._properties_ui.radioButton_execute_in_source.blockSignals(False)

    @Slot(str)
    def update_specification(self, text):
        """Update Tool specification according to selection in the specification comboBox.

        Args:
            text (str): Tool specification name in the comboBox
        """
        spec = self._project.get_specification(text)
        self.set_specification(spec)

    @Slot(bool)
    def _remove_arg(self, _=False):
        removed_rows = [index.row() for index in self._properties_ui.treeView_cmdline_args.selectedIndexes()]
        cmd_line_args = [arg for row, arg in enumerate(self.cmd_line_args) if row not in removed_rows]
        self._push_update_cmd_line_args_command(cmd_line_args)

    @Slot(bool)
    def _add_selected_file_path_args(self, _=False):
        new_args = [LabelArg(index.data()) for index in self._properties_ui.treeView_input_files.selectedIndexes()]
        self._push_update_cmd_line_args_command(self.cmd_line_args + new_args)

    @Slot(list)
    def _push_update_cmd_line_args_command(self, cmd_line_args):
        if self.cmd_line_args == cmd_line_args:
            return
        self._toolbox.undo_stack.push(UpdateCmdLineArgsCommand(self, cmd_line_args))

    def update_cmd_line_args(self, cmd_line_args):
        """Updates instance cmd line args list and sets the list as text to the line edit widget.

        Args:
            cmd_line_args (list): Tool cmd line args
        """
        self.cmd_line_args = cmd_line_args
        self._populate_cmdline_args_model()
        self._check_notifications()

    def _populate_cmdline_args_model(self):
        spec_args = [CmdLineArg(arg) for arg in self._specification.cmdline_args] if self._specification else []
        tool_args = self.cmd_line_args
        self._cmdline_args_model.reset_model(spec_args, tool_args)
        if self._active:
            self._properties_ui.treeView_cmdline_args.setFocus()

    def undo_specification(self):
        if self._specification is None:
            return None
        undo_spec = self._specification.clone()
        return undo_spec

    def do_set_specification(self, specification):
        """see base class"""
        if not super().do_set_specification(specification):
            return False
        self._populate_cmdline_args_model()
        if self._active:
            self._update_tool_ui()
        self._resources_to_successors_changed()
        self._check_notifications()
        return True

    def update_options(self, options):
        """Pushes a new UpdateToolOptionsCommand to the toolbox undo stack."""
        self._toolbox.undo_stack.push(UpdateToolOptionsCommand(self, options))

    def do_set_options(self, options):
        """Sets options for this tool.

        Args:
            options (dict): The new options dictionary, must include *ALL* the options, not only changed ones.
        """
        self._options = options
        if self._active:
            _ = self._get_options_widget()

    @Slot(bool)
    def _set_kill_completed_processes(self, kill_completed_processes):
        """Pushes a UpdateKillCompletedProcesses to the undo stack

        Args:
            kill_completed_processes (bool): whether to kill completed persistent processes after execution
        """
        self._toolbox.undo_stack.push(UpdateKillCompletedProcesses(self, kill_completed_processes))

    def do_update_kill_completed_processes(self, kill_completed_processes):
        """Updates the kill_completed_processes flag.

        Args:
            kill_completed_processes (bool): whether to kill completed persistent processes after execution
        """
        self._kill_completed_processes = kill_completed_processes
        if self._active:
            self._properties_ui.kill_consoles_check_box.setChecked(kill_completed_processes)

    @Slot(bool)
    def _set_log_process_output(self, log_process_output):
        """Pushes a UpdateLogProcessOutput to the undo stack

        Args:
            log_process_output (bool): whether to kill completed persistent processes after execution
        """
        self._toolbox.undo_stack.push(UpdateLogProcessOutput(self, log_process_output))

    def do_update_log_process_output(self, log_process_output):
        """Updates the log_process_output flag.

        Args:
            log_process_output (bool): whether to kill completed persistent processes after execution
        """
        self._log_process_output = log_process_output
        if self._active:
            self._properties_ui.log_process_output_check_box.setChecked(log_process_output)

    def _update_tool_ui(self):
        """Updates Tool properties UI. Used when Tool specification is changed.."""
        options_widget = self._properties_ui.horizontalLayout_options.takeAt(0)
        if options_widget:
            options_widget.widget().hide()
        if not self.specification():
            self._properties_ui.comboBox_tool.setCurrentIndex(-1)
            self.do_update_execution_mode(True)
            self._properties_ui.toolButton_tool_specification.setMenu(None)
            return
        self._properties_ui.comboBox_tool.setCurrentText(self.specification().name)
        self._update_specification_menu()
        self._properties_ui.toolButton_tool_specification.setMenu(self._specification_menu)
        options_widget = self._get_options_widget()
        if options_widget:
            self._properties_ui.horizontalLayout_options.addWidget(options_widget)
            options_widget.show()
        if self._specification.tooltype == "python":
            self.specification().set_execution_settings()
            k_spec_name = self.specification().execution_settings["kernel_spec_name"]
            env = self.specification().execution_settings["env"]
            use_console = self.specification().execution_settings["use_jupyter_console"]
            self._properties_ui.label_jupyter.show()
            if not use_console:
                exe = self.specification().execution_settings["executable"]
                p = resolve_python_interpreter(exe)
                self._properties_ui.label_jupyter.setText(f"[Basic console] {p}")
            else:
                env = "" if not env else f"[{env}]"
                self._properties_ui.label_jupyter.setText(f"[Jupyter Console] {k_spec_name} {env}")
        self._properties_ui.kill_consoles_check_box.setChecked(self._kill_completed_processes)
        self._update_kill_consoles_check_box_enabled()
        self._properties_ui.log_process_output_check_box.setChecked(self._log_process_output)

    def _update_specification_menu(self):
        spec_model_index = self._toolbox.specification_model.specification_index(self.specification().name)
        self._specification_menu = ToolSpecificationMenu(self._toolbox, spec_model_index, self)
        self._specification_menu.setTitle("Specification...")

    @Slot(bool)
    def _open_results_directory(self, _):
        """Open output directory in file browser."""
        if not os.path.exists(self.output_dir):
            self._logger.msg_warning.emit(f"Tool {self.name} has no results. Click Execute to generate them.")
            return
        url = "file:///" + self.output_dir
        # noinspection PyTypeChecker, PyCallByClass, PyArgumentList
        res = open_url(url)
        if not res:
            self._logger.msg_error.emit(f"Failed to open directory: {self.output_dir}")

    def specification(self):
        """Returns Tool specification."""
        return self._specification

    def resources_for_direct_successors(self):
        """See base class"""
        return scan_for_resources(self, self.specification(), self.output_dir)

    @property
    def executable_class(self):
        return ExecutableItem

    def _find_input_files(self, resources):
        """Iterates files in required input files model and looks for them in the given resources.

        Args:
            resources (list): resources available

        Returns:
            Dictionary mapping required files to path where they are found, or to None if not found
        """
        if not self.specification():
            return {}
        file_paths = dict()
        for req_file_path in self.specification().inputfiles:
            # Just get the filename if there is a path attached to the file
            _, filename = os.path.split(req_file_path)
            if not filename:
                # It's a directory
                continue
            file_paths[req_file_path] = find_file(filename, resources, one_file=True)
        return file_paths

    def _check_notifications(self):
        """See base class."""
        self.clear_notifications()
        self._check_write_index()
        if not self.specification():
            self.add_notification("This Tool does not have a specification. Set it in the Tool Properties Panel.")
        elif self.specification().includes and not self.specification().path:
            n = self.specification().name
            self.add_notification(
                f"Tool specification {n} path does not exist. Fix this in Tool specification editor."
            )
        resources = self._resources_from_upstream + self._resources_from_downstream
        resources_changed = resources != self._available_resources
        req_files_changed = False
        if self.specification():
            required_files = self.specification().inputfiles
            req_files_changed = required_files != self._req_file_paths
            # No need to perform slow check for if all the file requirements are/aren't still satisfied
            # if no changes have been made to the Tool's requirements or available resources
            if required_files and (req_files_changed or resources_changed):
                file_paths = self._find_input_files(resources)
                file_paths = flatten_file_path_duplicates(file_paths, self._logger)
                self._input_files_not_found = [k for k, v in file_paths.items() if v is None]
        if self._input_files_not_found:
            self.add_notification(
                "File(s) {0} needed to execute this Tool are not provided by any input item. "
                "Connect items that provide the required files to this Tool.".format(
                    ", ".join(self._input_files_not_found)
                )
            )
        if resources_changed or req_files_changed:
            # Only check for duplicate files in available resources when the resources have changed
            if resources_changed:
                self._available_resources = resources
                duplicates = self._input_file_model.duplicate_paths()
                if duplicates:
                    self.add_notification("Duplicate input files:
{}".format("
".join(duplicates)))
            if req_files_changed:
                self._req_file_paths = required_files
        missing_args = ", ".join(arg.arg for arg in self.cmd_line_args if isinstance(arg, LabelArg) and arg.missing)
        if missing_args:
            self.add_notification(
                f"The following command line argument(s) don't match any available resources: {missing_args}"
            )

    def handle_execution_successful(self, execution_direction, engine_state):
        """See base class."""
        super().handle_execution_successful(execution_direction, engine_state)
        if execution_direction != "FORWARD":
            return
        self._resources_to_successors_changed()

    def upstream_resources_updated(self, resources):
        """See base class."""
        self._resources_from_upstream = resources
        self._update_files_and_cmd_line_args()

    def replace_resources_from_upstream(self, old, new):
        """See base class."""
        self._replace_resources(old, new, self._resources_from_upstream)

    def downstream_resources_updated(self, resources):
        """See base class."""
        self._resources_from_downstream = resources
        self._update_files_and_cmd_line_args()

    def replace_resources_from_downstream(self, old, new):
        """See base class."""
        self._replace_resources(old, new, self._resources_from_downstream)

    def _replace_resources(self, old_resources, new_resources, resource_list):
        """Replaces resources.

        Modifies ``resource_list`` in-place!

        Args:
            old_resources (list of ProjectItemResource): old resources
            new_resources (list of ProjectItemResource): new resources
            resource_list (list of ProjectItemResource): current downstream or upstream resources
        """
        for old, new in zip(old_resources, new_resources):
            for i, resource in enumerate(resource_list):
                if resource == old:
                    resource_list[i] = new
                    break
            for i, arg in enumerate(self.cmd_line_args):
                if arg.arg == old.label:
                    self.cmd_line_args[i] = LabelArg(new.label)
        self._input_file_model.update(self._resources_from_upstream + self._resources_from_downstream)
        self.update_cmd_line_args(self.cmd_line_args)

    def _update_files_and_cmd_line_args(self):
        """Updates the file model and command line arguments."""
        resources = self._resources_from_upstream + self._resources_from_downstream
        self._input_file_model.update(resources)
        update_args = list()
        resource_labels = {resource.label for resource in resources}
        for arg in self.cmd_line_args:
            if arg.arg in resource_labels:
                arg.missing = False
            elif isinstance(arg, LabelArg):
                # Arg may refer to legacy Data connection resource labels that weren't prefixed by <project> or <data>
                legacy_arg_converted = False
                for label in resource_labels:
                    if label.startswith("<") and label.endswith(arg.arg):
                        legacy_arg_converted = True
                        arg.arg = label
                        break
                if not legacy_arg_converted:
                    arg.missing = True
            update_args.append(arg)
        self.update_cmd_line_args(update_args)

    def item_dict(self):
        """Returns a dictionary corresponding to this item."""
        d = super().item_dict()
        if not self.specification():
            d["specification"] = ""
        else:
            d["specification"] = self.specification().name
        d["execute_in_work"] = self.execute_in_work
        d["cmd_line_args"] = [arg.to_dict() for arg in self.cmd_line_args]
        if self._options:
            d["options"] = self._options
        d["kill_completed_processes"] = self._kill_completed_processes
        d["log_process_output"] = self._log_process_output
        if self._group_id:
            d["group_id"] = self._group_id
        return d

    @staticmethod
    def from_dict(name, item_dict, toolbox, project):
        """See base class."""
        description, x, y = DBWriterItemBase.parse_item_dict(item_dict)
        specification_name = item_dict.get("specification", "")
        execute_in_work = item_dict.get("execute_in_work", True)
        cmd_line_args = item_dict.get("cmd_line_args", [])
        cmd_line_args = [make_cmd_line_arg(arg) for arg in cmd_line_args]
        options = item_dict.get("options", {})
        kill_completed_processes = item_dict.get("kill_completed_processes", False)
        log_process_output = item_dict.get("log_process_output", False)
        group_id = item_dict.get("group_id")
        return Tool(
            name,
            description,
            x,
            y,
            toolbox,
            project,
            specification_name,
            execute_in_work,
            cmd_line_args,
            options,
            kill_completed_processes,
            log_process_output,
            group_id,
        )

    def rename(self, new_name, rename_data_dir_message):
        """See base class."""
        if not super().rename(new_name, rename_data_dir_message):
            return False
        self.output_dir = os.path.join(self.data_dir, TOOL_OUTPUT_DIR)
        return True

    def notify_destination(self, source_item):
        """See base class."""
        if source_item.item_type() == "Data Store":
            self._logger.msg.emit(
                f"Link established. Data Store {source_item.name} url will "
                f"be passed to Tool {self.name} when executing."
            )
        elif source_item.item_type() == "Data Connection":
            self._logger.msg.emit(
                f"Link established. Tool {self.name} will look for input "
                f"files from {source_item.name}'s references and data directory."
            )
        elif source_item.item_type() == "Exporter":
            self._logger.msg.emit(
                f"Link established. The file exported by {source_item.name} are now "
                f"available in {self.name}."
            )
        elif source_item.item_type() in ["Data Transformer", "Tool"]:
            self._logger.msg.emit("Link established")
        else:
            super().notify_destination(source_item)

    def actions(self):
        # pylint: disable=attribute-defined-outside-init
        if self.specification() is not None:
            if self._specification_menu is None:
                self._update_specification_menu()
            self._actions = [self._specification_menu.menuAction()]
        else:
            action = QAction("New specification")
            action.triggered.connect(self.show_specification_window)
            self._actions = [action]
        self._actions.append(QAction("Open results directory"))
        self._actions[-1].triggered.connect(self._open_results_directory)
        return self._actions

    @staticmethod
    def upgrade_v1_to_v2(item_name, item_dict):
        """Upgrades item's dictionary from v1 to v2.

        Changes:
        - 'tool' key is renamed to 'specification'

        Args:
            item_name (str): item's name
            item_dict (dict): Version 1 item dictionary

        Returns:
            dict: Version 2 Tool dictionary
        """
        item_dict["specification"] = item_dict.pop("tool", "")
        return item_dict


import inspect
import os
import sys
import threading
from html.parser import HTMLParser
import requests
from colorama import deinit
import pyjsparser
from lib.controller.controller import start, task_push_from_name
from lib.core.enums import HTTPMETHOD
from lib.parse.parse_request import FakeReq
from lib.parse.parse_responnse import FakeResp
from lib.proxy.baseproxy import AsyncMitmProxy
from html import unescape
import re
from html.parser import HTMLParser
from lib.helper.jscontext import analyse_js
from lib.parse.cmdparse import cmd_line_parser
from lib.core.data import logger, conf, KB
from lib.core.option import init
from abc import ABC
import copy
from lxml import etree
from jsonpath import jsonpath
import json
import string
import random
import argparse
from urllib.parse import urlparse,parse_qs
import time
class MyHTMLParser(HTMLParser, ABC):
    def __init__(self):
        super().__init__()
        self.tree = []
        self.tokenizer = []
        self.root = None
        temp = {
            "tagname": "",
            "content": "",
            "attibutes": []
        }

    def handle_starttag(self, tag, attrs):
        if len(self.tree) == 0:
            self.root = tag
        self.tree.append(
            {
                "tagname": tag,
                "content": "",
                "attibutes": attrs
            }
        )

    def handle_endtag(self, tag):
        if len(self.tree) > 0:
            r = self.tree.pop()
            self.tokenizer.append(r)

    def handle_startendtag(self, tag, attrs):
        self.handle_starttag(tag, attrs)
        self.handle_endtag(tag)

    def handle_data(self, data):
        if self.tree:
            self.tree[-1]["content"] += data

    def handle_comment(self, data):
        self.tokenizer.append({
            "tagname": "#comment",
            "content": data,
            "attibutes": []
        })

    def getTokenizer(self):
        while len(self.tree):
            r = self.tree.pop()
            self.tokenizer.append(r)
        return self.tokenizer

def getParamsFromHtml(html):
    parse = MyHTMLParser()
    parse.feed(html)
    tokens = parse.getTokenizer()

    result = set()
    for token in tokens:
        tagname = token["tagname"].lower()
        if tagname == "input":
            for attibute in token["attibutes"]:
                key, value = attibute
                if key == "name":
                    result.add(value)
                    break
        elif tagname == "script":
            content = token["content"]
            try:
                nodes = pyjsparser.parse(content).get("body", [])
            except pyjsparser.pyjsparserdata.JsSyntaxError as e:
                return []
            result |=set(analyse_js(nodes))
    return list(result)

def random_UA():
    ua_list = [
        'Mozilla/5.0 (Windows NT 6.1; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/30.0.1599.101',
        'Mozilla/5.0 (Windows NT 6.1; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/38.0.2125.122',
        'Mozilla/5.0 (Windows NT 6.1; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/39.0.2171.71',
        'Mozilla/5.0 (Windows NT 6.1; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/39.0.2171.95',
        'Mozilla/5.0 (Windows NT 6.1; WOW64) AppleWebKit/537.1 (KHTML, like Gecko) Chrome/21.0.1180.71',
        'Mozilla/4.0 (compatible; MSIE 6.0; Windows NT 5.1; SV1; QQDownload 732; .NET4.0C; .NET4.0E)',
        'Mozilla/5.0 (Windows NT 5.1; U; en; rv:1.8.1) Gecko/20061208 Firefox/2.0.0 Opera 9.50',
        'Mozilla/5.0 (Windows NT 6.1; WOW64; rv:34.0) Gecko/20100101 Firefox/34.0',
        'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/76.0.3809.100 '
        'Safari/537.36',
        'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_13_6) AppleWebKit/537.36 (KHTML, like Gecko) '
        'Chrome/76.0.3809.100 Safari/537.36',
        'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/76.0.3809.100 Safari/537.36',
        'Mozilla/5.0 (Windows NT 6.1; WOW64; rv:54.0) Gecko/20100101 Firefox/68.0',
        'Mozilla/5.0 (Macintosh; Intel Mac OS X 10.13; rv:61.0) Gecko/20100101 Firefox/68.0',
        'Mozilla/5.0 (X11; Linux i586; rv:31.0) Gecko/20100101 Firefox/68.0',
    ]
    return random.choice(ua_list)

def add_cmd_arg():
    parser = argparse.ArgumentParser(description='xss检测增加地址参数')
    parser.add_argument('-u',"--url",type=str,help='地址',default=None)
    parser.add_argument('-f',"--file_address",type=str,help='地址文件',default=None)
    args=parser.parse_args()
    return args
def get_able_xss_param(my_session,domain_url,data):#得到界面的有反射的参数
   # print("开始运行get_able_xss_param")
    able_param_inner=[]
    #data= {'exp': 'oooexpooo*', 'o': 'ooooooo*', 'l': 'ooolooo*', 'a': 'oooaooo*','exp1': 'oooexpooo*', 'o1': 'ooooooo*', 'l1': 'ooolooo*', 'a1': 'oooaooo*','l2': 'ooolooo*', 'a2': 'oooaooo*'}
    try:
        resp = my_session.get(domain_url,params=data, headers={"User-Agent":random_UA()}).text
    except Exception as e:
        print("网站无法访问！")
        return []
    for key,value in data.items():
        if value in resp:
            able_param_inner.append(key)
    return able_param_inner

def get_name_params(res):#得到input标签里的参数
    e=etree.HTML(res)
    name_params=set(e.xpath('//input/@name'))
    return name_params

def get__form_address(res):#得到form提交的地址
    e=etree.HTML(res)
    name_params=e.xpath('//form/@action')
    return name_params

def get_normal_xscan_params(s,url,res):#得到普通的参数
    params_10_number={}
    data={}
    params_data = {}
    parse_params={}
    able_xss_params=[]
    able_xss_url=[]

    #1 获取参数
    parse_params = set(getParamsFromHtml(res))
    RISK_PARAMS = {"id", 'action', 'type', 'm', 'callback', 'cb'} #常用高危参数
    domain_url=url.split("?")[0]
    query=urlparse(url).query#得到参数
    url_params=parse_qs(query)#转化为字典
    for key,walue in url_params.items():#取出参数
       url_params_key.append(key)
    parse_params = parse_params | RISK_PARAMS | set(url_params_key) - get_name_params(res)#得到常用参数和返回参数的并集
    print("普通参数是",parse_params)
   # 2.准备poc
    for key in parse_params:
        params_data[key] = "ooo"+key+"ooo*" #给每个参数配一个随机字符串
   # print(params_data)
    i=0
    for key,value in params_data.items():# 每10个参数放入字典中
       # print("开始放入字典")
        params_10_number[key]=value
        i=i+1
        if(len(params_10_number)==10 or i==len(params_data)):
            data=copy.deepcopy(params_10_number)
            params_10_number={}
            able_xss_params=able_xss_params+get_able_xss_param(s,domain_url,data)
            data={}
            time.sleep(2)

    for param in able_xss_params:
        able_xss_url.append(domain_url+"?"+param+"=ooo"+param+"ooo")
    print("能用的普通参数是",able_xss_params)
    print("能用的普通地址是",able_xss_url)
    return able_xss_url

def get_special_xscan_params(s,url,res):#得到form表格里的参数
    params_10_number={}
    data={}
    params_data = {}
    parse_params={}
    able_xss_params=[]
    form_address=[]
    new_address=[]
    able_xss_url=[]

    #1 获取参数
    domain_url = url.split("?")[0]
    parse_params = get_name_params(res)#得到input参数
    print("input参数是",parse_params)
    form_address=get__form_address(res)#得到form跳转地址
    print("form得到的地址是",form_address)
    domain_url_list=[]
    for address in form_address:
        if "http" in address:
            domain_url=address #如果form的地址是完整地址直接用http完整地址
            domain_url_list.append(domain_url)
        else:
           # print(str( domain_url.find("/")))
            if domain_url.count("/")>2:
                new_address=domain_url.split("/")
                domain_url=domain_url.replace("/"+new_address[len(new_address)-1],"")
                domain_url=domain_url+address
                domain_url_list.append(domain_url)
            else:
                domain_url=domain_url+address
                domain_url_list.append(domain_url)

   # 2.准备poc
    for key in parse_params:
            params_data[key] = "ooo"+key+"ooo*" #给每个参数配一个随机字符串
    for domain in domain_url_list:#取出一个提交地址
        print("form提交地址是",domain)
    # print(params_data)
        i=0
        for key,value in params_data.items():# 每4个参数放入字典中
           # print("放入的参数是",key)
            params_10_number[key]=value
            i=i+1
            if(len(params_10_number)==4 or i==len(params_data)):
                data=copy.deepcopy(params_10_number)
                params_10_number={}
                able_xss_params=able_xss_params+get_able_xss_param(s,domain,data)
                data={}
                time.sleep(2)

        for param in able_xss_params:
            able_xss_url.append(domain+"?"+param+"=ooo"+param+"ooo")
    print("能用的input参数是",able_xss_params)
    print("能用的input地址是",able_xss_url)
    return able_xss_url

def format_address(url):
    if url.find("//")==-1:#自动加http
        url="http://"+url
    if "?" in url:
        return url
    if(url[-1]=="/" or url[-1]=="\\"):#如果最后一位是斜杠
        url=url[:-1]
    return url

def possible_xss_url(s,urls):#判断页面是否有xss漏洞
    model1="@\"'<>@"
    model2="@\"'@"
    model3="@<>@"
    for url in urls:
        my_url=url.split('=')[0]
        if get_request_return(s,my_url,model1):
            xss_result=my_url+"="+model1+"存在xss反射漏洞！"
            return  xss_result
        else:
             if get_request_return(s,my_url,model2):
                xss_result=my_url+"="+model2+"存在xss反射漏洞！"
                return  xss_result
             else:
                if get_request_return(s,my_url,model3):
                    xss_result=my_url+"="+model3+"js上可能存在xss反射漏洞！"
                    return  xss_result
                else:
                    xss_result=""
                    return  xss_result
def get_request_return(s,url,model):#判断给的参数是否在返回的页面里
    my_url=url+"="+model
    try:
        resp = s.get(my_url, headers={"User-Agent":random_UA()}).text
        time.sleep(2)
    except Exception as e:
        print(e)
        return False
    #比较model与返回值是否相同
    if model in resp:
        return True
    else:
        return False
def test():
    url="http://demo.testfire.net"
    print(url.find('/'))

def main():
    global url_params_key
    url=""
    url_params_key=[]
    total_url=[]
    xss_result=""
    url_list=""
    total_xss_result=[]
    url=add_cmd_arg().url
    if url!=None:
        url_list.append(url)
    else:
        url_file_address=add_cmd_arg().file_address
        url_file_address=os.path.join(sys.path[0],'test.txt')
        print(url_file_address)
        with open(url_file_address, 'r',encoding='utf-8' ) as f:
            url_list=(f.read().strip()).split('\n')
    print(url_list)
    for url in url_list:
        print(url)
        url=url.strip()
        url=format_address(url)
        #print(url)
        try:
            res=requests.get(url).text
        except Exception as e:
            print(e)
            print("网站无法访问！")
            continue
        if(res==""):
            print("网站无法访问！")
            continue
        requests.DEFAULT_RETRIES=5
        s = requests.session()
        s.keep_alive=False
        total_url= list(set(get_normal_xscan_params(s,url,res)+get_special_xscan_params(s,url,res)))
        print(total_url)
        xss_result=possible_xss_url(s,total_url)
        if xss_result!="" and xss_result !=None:
            print(xss_result)
            total_xss_result.append(xss_result)

    print("总共找到如下地址xss",total_xss_result)

    #验证是否可能存在漏洞

    #resp = s.get(domain_url,params=params_data, headers={"User-Agent":random_UA()}).text

    #     #生成测试链接
    #     new_params=list(parse_params)
    #     for payload in payloads:
    #         for param in new_params:
    #            payload_url.append(url.split("?")[0]+"?"+param+"="+payload)
    # parse = MyHTMLParser()
    # parse.feed(t)
    # tokens = parse.getTokenizer()
    # print(tokens)
    # with open( 'test.txt', 'w',encoding='utf-8' ) as f:
    #     f.write(resp)

if __name__ == '__main__':
    main()


import os, json
import joblib
from model import LatentLinkerAgent

from utils import load_best_model, read_cli, infer_file
from environment.environments import Environment
from trainer import RLTrainer

args = read_cli()
print(args)
os.makedirs(args['dest_loc'], exist_ok=True)

if args['dataset'] == 'babi':
    from vocabulary.babi_agent_vocab import Vocabulary
elif args['dataset'] == 'bitod':
    from vocabulary.bitod_agent_vocab import Vocabulary
else:
    raise NotImplementedError

data_loc = args['data_loc']
vocab = joblib.load(os.path.join(args['dest_loc'], 'vocab.pkl'))

# 2. Get dataloaders
data_loc = args['data_loc']
batch_size = args['batch_size']
train_fname = os.path.join(data_loc, infer_file)

# 3. Get Environment
train_fname = os.path.join(args['data_loc'], infer_file)
eval_env = Environment(
    train_fname, vocab, mode='infer', batch_size=args['batch_size'],
    reward_model_dir=args['reward_model_loc'],
    device=args['device'], use_reward_table='train', track=True
)
args.update(vocab.get_ov_config(args))
model = load_best_model(args, LatentLinkerAgent)

trainer = RLTrainer(model, None, None, eval_env, device=args['device'])

ret, logs = trainer.eval_results()
print(json.dumps(ret, indent=2))
print(len(logs))

fname = os.path.join(args['data_loc'], infer_file[:-5] + '_pred.json')
with open(fname, 'w') as fp:
    json.dump(logs, fp, indent=2)


# unittest
import unittest
from readFile import readFile

def groupCrabs(data):
    result = {}
    for key in set(data):
        if key not in result:
            result[key] = 0
    for pos in data: result[pos] += 1
    return result

def costOfPosition(crabs, pos):
    keys = list(crabs.keys())
    cost = 0
    for key in keys:
        cost += abs(key - pos) * crabs[key]
    return cost

def costOfPositionWithIncrease(crabs, pos):
    keys = list(crabs.keys())
    cost = 0
    for key in keys:
        cost += sum(range(abs(key - pos)+1)) * crabs[key]
    return cost

def findCheapestCost(crabs, calcCost):
    cheapest_cost = float("inf")
    keys = list(crabs.keys())

    for i in range(keys[0],keys[-1]+1):
        cost_for_pos = calcCost(crabs, i)
        if cost_for_pos < cheapest_cost:
            cheapest_cost = cost_for_pos
    return cheapest_cost

class Test(unittest.TestCase):
    def testGroupCrabs(self):
        self.assertEqual(groupCrabs([16,1,2,1]), {16: 1, 1: 2, 2: 1})
        self.assertEqual(groupCrabs([16,1,2,0,4,2,7,1,2,14]), {0: 1, 1: 2, 2: 3, 4: 1, 7: 1, 14: 1, 16: 1})

    def testCostOfPosition(self):
        crabs = groupCrabs([16,1,2,0,4,2,7,1,2,14])
        result = costOfPosition(crabs, 2)
        self.assertEqual(result, 37)

    def testCostOfPositionWithIncrease(self):
        crabs = groupCrabs([16,1,2,0,4,2,7,1,2,14])
        result = costOfPositionWithIncrease(crabs, 5)
        self.assertEqual(result, 168)

    def testFindCheapestCost(self):
        crabs = groupCrabs([16,1,2,0,4,2,7,1,2,14])
        result = findCheapestCost(crabs, costOfPosition)
        self.assertEqual(result, 37)

    def testRun(self):
        raw_data = readFile('data')[0].split(',')
        data = [int(x) for x in raw_data]
        crabs = groupCrabs(data)
        self.assertEqual(findCheapestCost(crabs, costOfPosition), 326132)

        raw_data = readFile('data')[0].split(',')
        data = [int(x) for x in raw_data]
        crabs = groupCrabs(data)
        self.assertEqual(findCheapestCost(crabs, costOfPositionWithIncrease), 88612508)

if __name__ == '__main__':
    unittest.main()


'''
You are given an integer array arr. You can choose a set of integers and remove all the occurrences of these integers in the array.

Return the minimum size of the set so that at least half of the integers of the array are removed.

Example 1:

Input: arr = [3,3,3,3,5,5,5,2,2,7]
Output: 2
Explanation: Choosing {3,7} will make the new array [5,5,5,2,2] which has size 5 (i.e equal to half of the size of the old array).
Possible sets of size 2 are {3,5},{3,2},{5,2}.
Choosing set {2,7} is not possible as it will make the new array [3,3,3,3,5,5,5] which has a size greater than half of the size of the old array.
Example 2:

Input: arr = [7,7,7,7,7,7]
Output: 1
Explanation: The only possible set you can choose is {7}. This will make the new array empty.

Constraints:

2 <= arr.length <= 105
arr.length is even.
1 <= arr[i] <= 105
'''
from typing import *

import unittest
from collections import Counter

class Solution:
    def minSetSize(self, arr: List[int]) -> int:
        counter = Counter(arr)
        items = counter.items()
        items = sorted(items, key = lambda x: -x[1])
        target = len(arr) // 2
        count = 0
        for _, v in items:
            if target <= 0:
                return count
            target -= v
            count += 1
        return count

class TestSolution(unittest.TestCase):

    def test_case(self):
        examples = (
            (([9,77,63,22,92,9,14,54,8,38,18,19,38,68,58,19],), 5),
            (([3,3,3,3,5,5,5,2,2,7],), 2),
            (([7,7,7,7,7,7],), 1),
        )
        for first, second in examples:
            self.assert_function(first, second)

    def assert_function(self, first, second):
        self.assertEqual(Solution().minSetSize(*first), second,
                         msg="first: {}; second: {}".format(first, second))

unittest.main()

Leetcode/1338. Reduce Array Size to The Half.py


"""
给你两个整数数组 nums1 和 nums2 ，请你以数组形式返回两数组的交集。返回结果中每个元素出现的次数，
应与元素在两个数组中都出现的次数一致（如果出现次数不一致，则考虑取较小值）。可以不考虑输出结果的顺序。

示例 1：
输入：nums1 = [1,2,2,1], nums2 = [2,2]
输出：[2,2]

示例 2:
输入：nums1 = [4,9,5], nums2 = [9,4,9,8,4]
输出：[4,9]
"""

import collections
class Solution:
    #排序+双指针
    def intersection1(self,nums1,nums2):
        nums1.sort()
        nums2.sort()
        l1 = 0
        l2 = 0
        temp = []
        while l1<len(nums1) and l2<len(nums2):
            if nums1[l1] < nums2[l2]:
                l1 +=1
            elif nums1[l1] > nums2[l2]:
                l2 +=1
            else:
                temp.append(nums1[l1])
                l1 +=1
                l2 +=1
        return temp

    #哈希表
    def intersection(self, nums1, nums2):
        if len(nums1)>len(nums2):
            return self.intersection(nums2,nums1)
        # collections.Counter() 集合计数器
        m = collections.Counter(nums1)
        temp = []
        for i in range(len(nums2)):
            if nums2[i] in m and m[nums2[i]] > 0:
                temp.append(nums2[i])
                m[nums2[i]] -=1
        return temp

if __name__ == "__main__":
    nums1 = [1,2,2,1]
    nums2 = [2,2]
    s = Solution()
    res = s.intersection(nums1,nums2)
    print("res:",res)