import logging
import matplotlib.pyplot as plt
from dataclasses import dataclass

from eval_utils import parse_test_args, run_experiment, run_experiment_with_ui
from torchsim.core.datasets.dataset_se_base import SeDatasetSize
from torchsim.core.eval2.experiment import Experiment
from torchsim.core.eval2.experiment_controller import TrainTestComponentParams
from torchsim.core.eval2.experiment_runner_params import ExperimentParams
from torchsim.core.eval2.scaffolding import TopologyScaffoldingFactory
from torchsim.core.nodes import DatasetSeObjectsNode
from torchsim.research.experiment_templates2.task0_ta_analysis_template import Task0TaAnalysisTemplate, Task0TaAnalysisParams
from torchsim.research.research_topics.rt_2_1_2_learning_rate.node_groups.se_node_group import SeNodeGroup
from torchsim.research.research_topics.rt_3_7_1_task0_analysis.node_groups.dummy_model_group import DummyModelGroup
from torchsim.research.research_topics.rt_3_7_1_task0_analysis.topologies.task0_ta_analysis_topology import \
    Task0TaAnalysisTopology

logger = logging.getLogger(__name__)

@dataclass
class Params:
    experiment_params: Task0TaAnalysisParams
    train_test_params: TrainTestComponentParams

debug_params = Params(Task0TaAnalysisParams(measurement_period=1,
                                            num_classes=DatasetSeObjectsNode.label_size(),
                                            num_layers=2,  # has to be overwritten by actual num layers later
                                            sp_evaluation_period=2,
                                            show_conv_agreements=True),
                      TrainTestComponentParams(num_testing_phases=3,
                                               num_testing_steps=160,
                                               overall_training_steps=30))

middle_params = Params(Task0TaAnalysisParams(measurement_period=1,
                                             num_classes=DatasetSeObjectsNode.label_size(),
                                             num_layers=2,
                                             sp_evaluation_period=2),
                       TrainTestComponentParams(num_testing_phases=6,
                                                num_testing_steps=200,
                                                overall_training_steps=1500))

full_params = Params(Task0TaAnalysisParams(measurement_period=1,
                                           num_classes=DatasetSeObjectsNode.label_size(),
                                           num_layers=2,
                                           sp_evaluation_period=100),
                     TrainTestComponentParams(num_testing_phases=25,
                                              num_testing_steps=500,
                                              overall_training_steps=25000))

def run_measurement(name, topology_parameters, args, exp_pars):
    """"Runs the experiment with specified params, see the parse_test_args method for arguments"""

    scaffolding = TopologyScaffoldingFactory(Task0TaAnalysisTopology,
                                             se_group=SeNodeGroup,
                                             model=DummyModelGroup)

    template = Task0TaAnalysisTemplate("Task 0 layer-wise stats and classification accuracy",
                                       exp_pars.experiment_params,
                                       exp_pars.train_test_params)

    runner_parameters = ExperimentParams(max_steps=exp_pars.train_test_params.max_steps,
                                         save_cache=args.save,
                                         load_cache=args.load,
                                         clear_cache=args.clear,
                                         calculate_statistics=not args.computation_only,
                                         experiment_folder=args.alternative_results_folder)

    experiment = Experiment(template, scaffolding, topology_parameters, runner_parameters)

    logger.info(f'Running model: {name}')
    if args.run_gui:
        run_experiment_with_ui(experiment)
    else:
        run_experiment(experiment)

    if args.show_plots:
        plt.show()

def run_debug_comparison(args, num_conv_layers: int, exp_params):
    name = "Learning-rate-debug"

    cf_easy = [1, 2, 3, 4]
    size = SeDatasetSize.SIZE_24

    params = [
        {
            'se_group': {'class_filter': cf_easy,
                         'image_size': size},
            'model': {}
        },
        {
            'se_group': {'class_filter': cf_easy,
                         'image_size': size},
            'model': {}
        }
    ]

    exp_params.experiment_params.num_layers = num_conv_layers + 1  # needs to be there

    run_measurement(name, params, args, exp_params)

if __name__ == '__main__':
    arg = parse_test_args()

    run_debug_comparison(arg, 1, debug_params)


from __future__ import print_function

import sys
import time
from pyspark.sql import SparkSession

if __name__ == "__main__":
    """
        Usage: long_running [partitions] [run_time_sec]
    """
    partitions = int(sys.argv[1]) if len(sys.argv) > 1 else 5
    run_time_sec = int(sys.argv[2]) if len(sys.argv) > 2 else 600

    spark = SparkSession \
        .builder \
        .appName("Long-Running Spark Job") \
        .getOrCreate()

    n = 100000 * partitions
    data = spark.sparkContext.parallelize(range(1, n + 1), partitions)

    def processPartition(partition):
        """Sleep for run_time_sec"""
        print('Start processing partition')
        time.sleep(run_time_sec)
        print('Done processing partition')

    data.foreachPartition(processPartition)
    print('Job completed successfully')

    spark.stop()


import os
import random
import asyncio
from io import BytesIO

import aiohttp
from graia.saya import Saya, Channel
from graia.ariadne.model import Group, Member
from graia.ariadne.message.element import Image
from graia.broadcast.interrupt.waiter import Waiter
from graia.broadcast.interrupt import InterruptControl
from graia.ariadne.event.message import GroupMessage
from graia.ariadne.message.chain import MessageChain
from graia.saya.builtins.broadcast.schema import ListenerSchema
from graia.ariadne.message.parser.twilight import Twilight, RegexMatch, RegexResult

from util.sendMessage import safeSendGroupMessage
from util.control import Permission, Interval
from .WaifuLab import waifu

func = os.path.dirname(__file__).split("\\")[-1]

saya = Saya.current()
channel = Channel.current()
channel.name(func)
bcc = saya.broadcast
inc = InterruptControl(bcc)

http_proxy = {"http://127.0.0.1:33210"}

BROWER = 0

async def get_pic(pic_id: str):
    url = f"https://www.thiswaifudoesnotexist.net/example-{pic_id}.jpg"
    async with aiohttp.ClientSession() as session:
        resp = await session.get(url=url)
        res = await resp.content.read()

    return BytesIO(res).getvalue()

@channel.use(
    ListenerSchema(
        listening_events=[GroupMessage],
        inline_dispatchers=[
            Twilight(
                [
                    RegexMatch("/"),
                    "head" @ RegexMatch("Waifu|头像生成|获取头像|随机头像"),
                ]
            )
        ],
        decorators=[
            Permission.require(),
            Permission.restricter(func),
            Interval.require(),
        ],
    )
)
async def main(group: Group):
    ran = random.randint(1, 100000)
    img = await get_pic(ran)
    await safeSendGroupMessage(
        group, MessageChain.create("你需要的头像已经生成好了哦~", Image(data_bytes=img))
    )

@channel.use(
    ListenerSchema(
        listening_events=[GroupMessage],
        inline_dispatchers=[
            Twilight(
                [
                    RegexMatch("/"),
                    "head" @ RegexMatch("WaifuLab|Lab头像生成|Lab获取头像|Lab头像"),
                    "water" @ RegexMatch("无水印", optional=True),
                ]
            )
        ],
        decorators=[
            Permission.require(),
            Permission.restricter(func),
            Interval.require(),
        ],
    )
)
async def WaifuLab(group: Group, member: Member, water: RegexResult):
    global BROWER
    if BROWER > 5:
        return await safeSendGroupMessage(group, MessageChain.create("当前进程过多，等一会再来哦~"))
    else:
        BROWER += 1
    if water.matched:
        page = waifu(water_mark=False)
    else:
        page = waifu()
    await safeSendGroupMessage(
        group,
        "开始获取头像了哦~\n请耐心等待~",
    )
    stat = await page.gey_browser()
    if not stat:
        BROWER -= 1
        return await safeSendGroupMessage(group, MessageChain.create("连接超时，等一会再来哦~"))
    pic = await page.first_shot()
    pic = pic.getvalue()
    await safeSendGroupMessage(
        group,
        MessageChain.create("连接成功，请选择下一步哦~\n选择[序号]\n继续\n返回\n退出", Image(data_bytes=pic)),
    )
    func_dict = {
        "选择": page.other_shot,
        "继续": page.continues,
        "返回": page.back,
        "退出": page.close,
    }

    @Waiter.create_using_function(
        listening_events=[GroupMessage],
        using_dispatchers=[
            Twilight(
                [
                    "head" @ RegexMatch("选择|继续|返回|退出"),
                    "choose" @ RegexMatch("[0-9]+", optional=True),
                ]
            )
        ],
    )
    async def get_waifu(waiter_member: Member, head: RegexResult, choose: RegexResult):
        tar = 0
        if member.id == waiter_member.id:
            head = head.result.asDisplay()
            if choose.matched:
                choose = int(choose.result.asDisplay())
                if choose > 16:
                    await safeSendGroupMessage(group, "超出选择范围了哦~")
                    return
                if choose == 0:
                    await safeSendGroupMessage(group, "没有0这个选项哦~")
                    return
                tar = choose
            await safeSendGroupMessage(
                group,
                "操作成功~请静待结果~",
            )
            img = func_dict[head]
            if tar != 0:
                img = await img(choose)
            else:
                img = await img()
            if not img:
                return False  # 当返回为空时，选项为退出
            elif isinstance(img, tuple):
                await safeSendGroupMessage(
                    group,
                    MessageChain.create(
                        "获取成功~",
                        Image(data_bytes=img[0].getvalue()),
                        Image(data_bytes=img[1].getvalue()),
                    ),
                )
                return True
            else:
                await safeSendGroupMessage(
                    group,
                    MessageChain.create("头像获取成功~", Image(data_bytes=img.getvalue())),
                )
                return False
        return

    sign = True
    while sign:
        try:
            sign = await asyncio.wait_for(inc.wait(get_waifu), 30)
        except asyncio.TimeoutError:
            await page.close()
            BROWER -= 1
            return await safeSendGroupMessage(group, "等待超时~")
    BROWER -= 1

    await safeSendGroupMessage(group, "你需要的头像已经生成好了哦~")


import requests
from faker import Faker

URL = "http://localhost:8080"

def ping():
    r = requests.get(URL + "/ping")
    assert r.status_code == 200
    assert r.json() == {"message": "pong"}

def gen_location():
    points = [Faker().latlng() for _ in range(3)]
    points = list(map(lambda x: [float(x[0]), float(x[1])], points))
    location = {
        "name": Faker().word(),
        "description": Faker().sentence(),
        "points": points,
    }
    return dict(location)

def make_location():
    location = gen_location()
    r = requests.post(URL + "/location", json=location)
    print(r.json())
    assert r.status_code == 201
    return r.json()

def get_location(location_id):
    r = requests.get(URL + f"/location/{location_id}")
    assert r.status_code == 200
    return r.json()

def delete_location(location_id):
    r = requests.delete(URL + f"/location/{location_id}")
    assert r.status_code == 200

def get_locations():
    r = requests.get(URL + "/locations")
    assert r.status_code == 200
    return r.json()

def get_nearby_locations(lat, lon, radius):
    r = requests.get(URL + f"/locations/nearby?lat={lat}&lon={lon}&radius={radius}")
    assert r.status_code == 200
    return r.json()

def get_images(location_id):
    r = requests.get(URL + f"/location/{location_id}/images")
    assert r.status_code == 200
    return r.json()

def get_image(image_id):
    r = requests.get(URL + f"/image/{image_id}")
    assert r.status_code == 200
    content = r.content
    return content

def upload_image(location_id):
    r = requests.post(
        URL + f"/upload_image/{location_id}", files={"file": open("image.png", "rb")}
    )
    assert r.status_code == 201
    return r.json()

def delete_image(image_id):
    r = requests.delete(URL + f"/image/{image_id}")
    assert r.status_code == 200

def get_image_data(image_id):
    r = requests.get(URL + f"/image/{image_id}/data")
    assert r.status_code == 200
    return r.json()

def create_user():
    user = {
        "username": Faker().user_name(),
        "name": Faker().name(),
        "profile_picture": Faker().image_url(),
    }
    r = requests.post(URL + "/user", json=user)
    assert r.status_code == 201
    return r.json()

def get_user(user_id):
    r = requests.get(URL + f"/user/{user_id}")
    assert r.status_code == 200
    return r.json()

def get_user_visits(user_id):
    r = requests.get(URL + f"/user/{user_id}/visits")
    assert r.status_code == 200
    return r.json()

def create_visit(user_id, location_id):
    r = requests.post(URL + f"/user/{user_id}/visit/{location_id}")
    assert r.status_code == 201
    return r.json()

if __name__ == "__main__":
    print("Checking server is up")
    try:
        ping()
    except Exception as e:
        print("Server is down")
        raise e

    print("Test make location")
    location = make_location()
    location_id = location["id"]
    print(location)

    print("Test get location")
    location = get_location(location_id)

    print("Test get locations")
    locations = get_locations()
    assert location in locations

    print("Test upload image")
    image = upload_image(location_id)
    image_id = image["id"]

    print("Test get image")
    image = get_image(image_id)
    assert image == open("image.png", "rb").read()

    print("Test get image data")
    image_data = get_image_data(image_id)
    assert image_data["location_id"] == location_id

    print("Test get images")
    images = get_images(location_id)
    assert image_data in images

    print("Test get nearby")
    nearby = get_nearby_locations(
        location["geometry"]["coordinates"][0][0][0],
        location["geometry"]["coordinates"][0][0][1],
        1000,
    )
    assert location in nearby

    print("Test create user")
    user = create_user()
    user_id = user["username"]
    assert user == get_user(user_id)

    print("Test create visit")
    visit = create_visit(user_id, location_id)
    assert visit in get_user_visits(user_id)


import warnings
import gymnasium as gym
import torch as t
import wandb
import time

from argparse import Namespace
from typing import Optional

from src.config import RunConfig, TransformerModelConfig, EnvironmentConfig, OnlineTrainConfig
from src.ppo.utils import set_global_seeds
from src.ppo.train import train_ppo
from src.utils import TrajectoryWriter
from src.environments.environments import make_env
from src.environments.registration import register_envs

warnings.filterwarnings("ignore", category=DeprecationWarning)
device = t.device("cuda" if t.cuda.is_available() else "cpu")

def ppo_runner(
    run_config: RunConfig,
    environment_config: EnvironmentConfig,
    online_config: OnlineTrainConfig,
    transformer_model_config: Optional[TransformerModelConfig],
):
    '''
    Trains a Proximal Policy Optimization (PPO) algorithm on a specified environment using the given hyperparameters.

    Args:
    - args (PPOArgs): an object that contains hyperparameters and other arguments for PPO training.

    Returns: None.
    '''

    if online_config.trajectory_path:
        trajectory_writer = TrajectoryWriter(
            online_config.trajectory_path,
            run_config=run_config,
            environment_config=environment_config,
            online_config=online_config,
            transformer_model_config=transformer_model_config,
        )
    else:
        trajectory_writer = None

    # Verify environment is registered
    register_envs()
    all_envs = [env_spec for env_spec in gym.envs.registry]
    assert environment_config.env_id in all_envs, f"Environment {environment_config.env_id} not registered."

    # wandb initialisation,
    run_name = f"{environment_config.env_id}__{run_config.exp_name}__{run_config.seed}__{int(time.time())}"
    if run_config.track:
        run = wandb.init(
            project=run_config.wandb_project_name,
            entity=run_config.wandb_entity,
            config=combine_args(
                run_config, environment_config, online_config, transformer_model_config
            ),  # vars is equivalent to args.__dict__
            name=run_name,
            save_code=True,
        )

    # add run_name to args
    run_config.run_name = run_name

    # make envs
    set_global_seeds(run_config.seed)

    envs = gym.vector.SyncVectorEnv(
        [make_env(
            env_id=environment_config.env_id,
            seed=environment_config.seed + i,
            idx=i,
            capture_video=environment_config.capture_video,
            run_name=run_name,
            max_steps=environment_config.max_steps,
            fully_observed=environment_config.fully_observed,
            flat_one_hot=environment_config.one_hot_obs,
            agent_view_size=environment_config.view_size,
            render_mode="rgb_array",
        ) for i in range(online_config.num_envs)]
    )

    agent = train_ppo(
        run_config=run_config,
        online_config=online_config,
        environment_config=environment_config,
        transformer_model_config=transformer_model_config,
        envs=envs,
        trajectory_writer=trajectory_writer,
        probe_idx=None
    )
    if run_config.track:
        run.finish()

def combine_args(run_config, environment_config, online_config, transformer_model_config=None):
    args = {}
    args.update(run_config.__dict__)
    args.update(environment_config.__dict__)
    args.update(online_config.__dict__)
    if transformer_model_config is not None:
        args.update(transformer_model_config.__dict__)
    return args


from django.shortcuts import redirect

from accounts.constants import *
from accounts.forms import UpdateUserBiographyForm
from accounts.forms import UpdateUserCityForm
from accounts.forms import UpdateUserFacebookForm
from accounts.forms import UpdateUserFirstNameForm
from accounts.forms import UpdateUserLastNameForm
from accounts.forms import UpdateUserShareLocationForm
from accounts.forms import UpdateUserShareNameForm
from accounts.forms import UpdateUserStateForm
from accounts.forms import UpdateUserTwitterForm
from accounts.forms import UpdateUserWebsiteForm
from accounts.forms import UpdateUsernameForm
from accounts.session_keys import *

from hacktoolkit.view_helpers import wrap_data

def wrap_data_accounts(request, data = None):
    if data == None:
        data = {}
    data = wrap_data(request, data)

    # logged in
    user = data.get('user', None)
    if user:
        pass
    else:
        pass
    return data

def get_social_auths_status(user):
    user_social_auths = user.profile.get_social_auths()
    status_dict = {}
    for social_auth in user_social_auths:
        #status_dict[social_auth.provider] = social_auth.id
        status_dict[social_auth.provider] = True

    status_list = []
    for social_auth in SOCIAL_AUTHS:
        key = social_auth['key']
        name = social_auth['name']
        item = {
            'key' : key,
            'name' : name,
            #'id' : status_dict.get(key, None),
            'linked' : status_dict.get(key, False),
        }
        status_list.append(item)
    return status_list

def redirect_to_social_auth_complete(request):
    """Return an HTTP Redirect response to social:complete to continue the pipeline
    """
    backend = request.session[SOCIAL_AUTH_PARTIAL_PIPELINE_KEY]['backend']
    response = redirect('social:complete', backend=backend)
    return response

def get_social_auths_status(user):
    user_social_auths = user.profile.get_social_auths()
    status_dict = {}
    for social_auth in user_social_auths:
        status_dict[social_auth.provider] = True

    status_list = []
    for social_auth in SOCIAL_AUTHS:
        key = social_auth['key']
        name = social_auth['name']
        item = {
            'key' : key,
            'name' : name,
            'linked' : status_dict.get(key, False)
        }
        status_list.append(item)
    return status_list

def get_user_update_forms(user):
    user_update_forms = [
        {
            'key' : 'update_username_form',
            'form' : UpdateUsernameForm(user, initial = {'username' : user.username if user.profile.has_username_set else '' }),
        },
        {
            'key' : 'update_user_first_name_form',
            'form' : UpdateUserFirstNameForm(user, instance=user),
        },
        {
            'key' : 'update_user_last_name_form',
            'form' : UpdateUserLastNameForm(user, instance=user),
        },
        {
            'key' : 'update_user_share_name_form',
            'form' : UpdateUserShareNameForm(user, instance=user.profile),
        },
        {
            'key' : 'update_user_website_form',
            'form' : UpdateUserWebsiteForm(user, instance=user.profile),
        },
        {
            'key' : 'update_user_facebook_form',
            'form' : UpdateUserFacebookForm(user, instance=user.profile),
        },
        {
            'key' : 'update_user_twitter_form',
            'form' : UpdateUserTwitterForm(user, instance=user.profile),
        },
        {
            'key' : 'update_user_city_form',
            'form' : UpdateUserCityForm(user, instance=user.profile),
        },
        {
            'key' : 'update_user_state_form',
            'form' : UpdateUserStateForm(user, instance=user.profile),
        },
        {
            'key' : 'update_user_share_location_form',
            'form' : UpdateUserShareLocationForm(user, instance=user.profile),
        },
        {
            'key' : 'update_user_biography_form',
            'form' : UpdateUserBiographyForm(user, instance=user.profile),
        },
    ]
    return user_update_forms


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
This code will convert the dataset's pixel values to Images
@author: raz
"""

import os
import csv
import argparse
import numpy as np
import scipy.misc

parser = argparse.ArgumentParser()
parser.add_argument('-f', '--file', required=True, help="path of the csv file")
parser.add_argument('-o', '--output', required=True, help="path of the output directory")
args = parser.parse_args()

w, h = 48, 48
image = np.zeros((h, w), dtype=np.uint8)
id = 1

with open(args.file) as csvfile:
    datareader = csv.reader(csvfile, delimiter =',')
    next(datareader,None)

    for row in datareader:

        emotion = row[0]
        pixels = row[1].split()
        usage = row[2]
        pixels_array = np.asarray(pixels, dtype=np.int)

        image = pixels_array.reshape(w, h)
        stacked_image = np.dstack((image,) * 3)

        image_folder = os.path.join(args.output, usage)
        if not os.path.exists(image_folder):
            os.makedirs(image_folder)
        image_file =  os.path.join(image_folder , emotion +'_'+ str(id) +'.jpg')
        scipy.misc.imsave(image_file, stacked_image)
        id += 1
        if id % 100 == 0:
            print('Processed {} images'.format(id))

print("Finished conversion to {} images".format(id))


import os
from unittest.mock import patch

import numpy as np
import pytest
import rasterio
import requests
from constants import KERNEL_SIZE
from shapely.geometry import Polygon
from write_to_db import (  # to_datetime_str,
    calc_peak_power,
    create_mask,
    create_polygon_and_area,
    extract_polygon_coordinates,
    extract_shapes,
    prediction_to_polygons,
    transform_polygon,
    write_to_db,
)

def test_create_mask_with_all_values_below_threshold():
    # Test with a prediction where all values are below the threshold
    prediction = np.random.rand(1, KERNEL_SIZE, KERNEL_SIZE) * 0.4
    expected = np.zeros((KERNEL_SIZE, KERNEL_SIZE), dtype=np.uint8)
    result = create_mask(prediction)
    assert result.shape == (KERNEL_SIZE, KERNEL_SIZE)
    np.testing.assert_array_equal(result, expected)

def test_create_mask_with_all_values_above_threshold():
    # Test with a prediction where all values are above the threshold
    prediction = np.random.rand(1, KERNEL_SIZE, KERNEL_SIZE) * 0.6 + 0.5
    expected = np.ones((KERNEL_SIZE, KERNEL_SIZE), dtype=np.uint8)
    result = create_mask(prediction)
    assert result.shape == (KERNEL_SIZE, KERNEL_SIZE)
    np.testing.assert_array_equal(result, expected)

def test_create_mask_with_some_values_above_threshold():
    # Test with a prediction where some values are above and below the threshold
    prediction = np.zeros((1, KERNEL_SIZE, KERNEL_SIZE))
    prediction[
        0, : KERNEL_SIZE // 2, : KERNEL_SIZE // 2
    ] = 0.4  # values below the threshold
    prediction[
        0, KERNEL_SIZE // 2 :, KERNEL_SIZE // 2 :
    ] = 0.6  # values above the threshold
    expected = np.zeros((KERNEL_SIZE, KERNEL_SIZE), dtype=np.uint8)
    expected[KERNEL_SIZE // 2 :, KERNEL_SIZE // 2 :] = 1
    result = create_mask(prediction)
    assert result.shape == (KERNEL_SIZE, KERNEL_SIZE)
    np.testing.assert_array_equal(result, expected)

def test_create_mask_with_invalid_input():
    # Test with invalid input
    prediction = "not a numpy array"
    with pytest.raises(TypeError):
        create_mask(prediction)

def test_create_mask_with_invalid_shape():
    # Test with invalid shape
    prediction = np.random.rand(2, KERNEL_SIZE, KERNEL_SIZE)
    with pytest.raises(ValueError):
        create_mask(prediction)

def test_create_mask_with_invalid_dtype():
    # Test with invalid dtype
    prediction = np.random.randint(0, 2, (1, KERNEL_SIZE, KERNEL_SIZE))
    with pytest.raises(TypeError):
        create_mask(prediction)

def test_extract_shapes_with_invalid_dtype():
    # Test with invalid dtype
    mask = np.random.rand(5, 5)
    transform = rasterio.transform.from_origin(0, 5, 1, 1)
    with pytest.raises(TypeError):
        extract_shapes(mask, transform)

def test_extract_shapes_with_invalid_dimensions():
    # Test with invalid dimensions
    mask = np.random.randint(0, 2, (1, 5, 5), dtype=np.uint8)
    transform = rasterio.transform.from_origin(0, 5, 1, 1)
    with pytest.raises(ValueError):
        extract_shapes(mask, transform)

def test_extract_shapes_with_single_shape():
    # Test with a mask that contains a single shape
    mask = np.zeros((5, 5), dtype=np.uint8)
    mask[1:4, 1:4] = 1
    transform = rasterio.transform.from_origin(0, 5, 1, 1)
    expected = [
        (
            {
                "type": "Polygon",
                "coordinates": [
                    [(1.0, 4.0), (1.0, 1.0), (4.0, 1.0), (4.0, 4.0), (1.0, 4.0)]
                ],
            },
            1.0,
        )
    ]
    result = extract_shapes(mask, transform)
    assert len(result) == 1
    assert result == expected

def test_extract_shapes_with_multiple_shapes():
    # Test with a mask that contains multiple shapes
    mask = np.zeros((5, 5), dtype=np.uint8)
    mask[1:3, 1:3] = 1
    mask[3:5, 3:5] = 1
    transform = rasterio.transform.from_origin(0, 5, 1, 1)
    expected = [
        (
            {
                "type": "Polygon",
                "coordinates": [
                    [(1.0, 4.0), (1.0, 2.0), (3.0, 2.0), (3.0, 4.0), (1.0, 4.0)]
                ],
            },
            1.0,
        ),
        (
            {
                "type": "Polygon",
                "coordinates": [
                    [(3.0, 2.0), (3.0, 0.0), (5.0, 0.0), (5.0, 2.0), (3.0, 2.0)]
                ],
            },
            1.0,
        ),
    ]
    result = extract_shapes(mask, transform)
    assert len(result) == 2
    assert result == expected
    assert isinstance(result, list)

def test_extract_shapes_with_no_shapes():
    # Test with a mask that contains no shapes
    mask = np.zeros((5, 5), dtype=np.uint8)
    transform = rasterio.transform.from_origin(0, 5, 1, 1)
    result = extract_shapes(mask, transform)
    assert len(result) == 0

def test_create_polygon_and_area_with_valid_input():
    # Test with valid input
    shape = list(
        (
            {
                "type": "Polygon",
                "coordinates": [
                    [(1.0, 4.0), (1.0, 1.0), (4.0, 1.0), (4.0, 4.0), (1.0, 4.0)]
                ],
            },
            1.0,
        )
    )
    expected_polygon = Polygon(
        [(1.0, 4.0), (1.0, 1.0), (4.0, 1.0), (4.0, 4.0), (1.0, 4.0)]
    )
    expected_area = expected_polygon.area
    result_polygon, result_area = create_polygon_and_area(shape)
    assert result_polygon.equals(expected_polygon)
    assert result_area == expected_area

def test_create_polygon_and_area_with_invalid_input():
    # Test with invalid input
    shape = "not a valid shape"
    with pytest.raises(TypeError):
        create_polygon_and_area(shape)

def test_create_polygon_and_area_with_invalid_coordinates():
    # Test with invalid coordinates
    shape = list(
        (
            {
                "type": "Polygon",
                "coordinates": "not a valid coordinates",
            },
            1.0,
        )
    )
    with pytest.raises(ValueError):
        create_polygon_and_area(shape)

def test_transform_polygon_with_valid_input():
    # Test with valid input
    polygon = Polygon([(1.0, 4.0), (1.0, 1.0), (4.0, 1.0), (4.0, 4.0), (1.0, 4.0)])
    crs = "EPSG:3857"
    result = transform_polygon(polygon, crs)
    assert isinstance(result, Polygon)
    assert result.is_valid

def test_transform_polygon_with_invalid_polygon():
    # Test with invalid polygon
    polygon = "not a valid polygon"
    crs = "EPSG:3857"
    with pytest.raises(TypeError):
        transform_polygon(polygon, crs)

def test_transform_polygon_with_invalid_crs():
    # Test with invalid CRS
    polygon = Polygon([(1.0, 4.0), (1.0, 1.0), (4.0, 1.0), (4.0, 4.0), (1.0, 4.0)])
    crs = 1234
    with pytest.raises(TypeError):
        transform_polygon(polygon, crs)

def test_transform_polygon_with_nonexistent_crs():
    # Test with non-existent CRS
    polygon = Polygon([(1.0, 4.0), (1.0, 1.0), (4.0, 1.0), (4.0, 4.0), (1.0, 4.0)])
    crs = "EPSG:9999"
    with pytest.raises(ValueError):
        transform_polygon(polygon, crs)

def test_prediction_to_polygons_with_valid_input():
    # Test with valid input
    prediction = np.random.rand(1, KERNEL_SIZE, KERNEL_SIZE)
    metadata = {
        "transform": rasterio.transform.from_origin(0, 5, 1, 1),
        "crs": "EPSG:3857",
    }
    polygons, areas = prediction_to_polygons(prediction, metadata)
    assert isinstance(polygons, list)
    assert isinstance(areas, list)
    assert len(polygons) == len(areas)

def test_prediction_to_polygons_with_invalid_prediction():
    # Test with invalid prediction
    prediction = "not a numpy array"
    metadata = {
        "transform": rasterio.transform.from_origin(0, 5, 1, 1),
        "crs": "EPSG:3857",
    }
    with pytest.raises(TypeError):
        prediction_to_polygons(prediction, metadata)

def test_prediction_to_polygons_with_invalid_metadata():
    # Test with invalid metadata
    prediction = np.random.rand(1, KERNEL_SIZE, KERNEL_SIZE)
    metadata = "not a dictionary"
    with pytest.raises(TypeError):
        prediction_to_polygons(prediction, metadata)

def test_prediction_to_polygons_with_missing_transform_in_metadata():
    # Test with missing 'transform' key in metadata
    prediction = np.random.rand(1, KERNEL_SIZE, KERNEL_SIZE)
    metadata = {"crs": "EPSG:3857"}
    with pytest.raises(KeyError):
        prediction_to_polygons(prediction, metadata)

def test_prediction_to_polygons_with_missing_crs_in_metadata():
    # Test with missing 'crs' key in metadata
    prediction = np.random.rand(1, KERNEL_SIZE, KERNEL_SIZE)
    metadata = {"transform": rasterio.transform.from_origin(0, 5, 1, 1)}
    with pytest.raises(KeyError):
        prediction_to_polygons(prediction, metadata)

# def test_to_datetime_str_with_valid_input():
#     # Test with valid input
#     date_string = "2022-01-01"
#     expected = "2022-01-01"
#     result = to_datetime_str(date_string)
#     assert result == expected

# def test_to_datetime_str_with_invalid_format():
#     # Test with date string in invalid format
#     date_string = "01-01-2022"
#     with pytest.raises(ValueError):
#         to_datetime_str(date_string)

# def test_to_datetime_str_with_invalid_type():
#     # Test with non-string input
#     date_string = 20220101
#     with pytest.raises(TypeError):
#         to_datetime_str(date_string)

# def test_to_datetime_str_with_empty_string():
#     # Test with empty string
#     date_string = ""
#     with pytest.raises(ValueError):
#         to_datetime_str(date_string)

def test_extract_polygon_coordinates_with_valid_input():
    # Test with valid input
    polygon = Polygon([(0, 0), (1, 1), (1, 0)])
    expected_latitudes = [0, 1, 0, 0]
    expected_longitudes = [0, 1, 1, 0]
    result_latitudes, result_longitudes = extract_polygon_coordinates(polygon)
    assert result_latitudes == expected_latitudes
    assert result_longitudes == expected_longitudes

def test_extract_polygon_coordinates_with_invalid_type():
    # Test with non-Polygon input
    polygon = "not a polygon"
    with pytest.raises(TypeError):
        extract_polygon_coordinates(polygon)

def test_extract_polygon_coordinates_with_no_exterior_ring():
    # Test with Polygon that has no exterior ring
    polygon = Polygon()
    with pytest.raises(ValueError):
        extract_polygon_coordinates(polygon)

def test_calc_peak_power_with_valid_input():
    # Test with valid input
    area_in_sq_m = 10000
    expected = 1.0
    result = calc_peak_power(area_in_sq_m)
    assert result == expected

def test_calc_peak_power_with_zero():
    # Test with zero area
    area_in_sq_m = 0
    expected = 0.0
    result = calc_peak_power(area_in_sq_m)
    assert result == expected

def test_calc_peak_power_with_negative_input():
    # Test with negative area
    area_in_sq_m = -1
    with pytest.raises(ValueError):
        calc_peak_power(area_in_sq_m)

def test_calc_peak_power_with_non_number_input():
    # Test with non-number input
    area_in_sq_m = "not a number"
    with pytest.raises(TypeError):
        calc_peak_power(area_in_sq_m)

@patch("requests.post")
@patch.dict(os.environ, {"JWT": "test_token"}, clear=True)
def test_write_to_db_with_valid_input(mock_post):
    # Test with valid input
    mock_post.return_value.raise_for_status.return_value = None
    polygon = Polygon([(0, 0), (1, 1), (1, 0)])
    area = 0.5
    tile_date = "2022-01-01"
    filename = "test.tif"
    assert write_to_db(polygon, area, tile_date, filename)

@patch("requests.post")
@patch.dict(os.environ, {"JWT": "test_token"}, clear=True)
def test_write_to_db_with_invalid_polygon(mock_post):
    # Test with invalid polygon
    polygon = "not a polygon"
    area = 0.5
    tile_date = "2022-01-01"
    filename = "test.tif"
    with pytest.raises(TypeError):
        write_to_db(polygon, area, tile_date, filename)

@patch("requests.post")
@patch.dict(os.environ, {"JWT": "test_token"}, clear=True)
def test_write_to_db_with_invalid_area(mock_post):
    # Test with invalid area
    polygon = Polygon([(0, 0), (1, 1), (1, 0)])
    area = -1
    tile_date = "2022-01-01"
    filename = "test.tif"
    with pytest.raises(ValueError):
        write_to_db(polygon, area, tile_date, filename)

@patch("requests.post")
@patch.dict(os.environ, {"JWT": "test_token"}, clear=True)
def test_write_to_db_with_invalid_tile_date(mock_post):
    # Test with invalid tile_date
    polygon = Polygon([(0, 0), (1, 1), (1, 0)])
    area = 0.5
    tile_date = 20220101
    filename = "test.tif"
    with pytest.raises(TypeError):
        write_to_db(polygon, area, tile_date, filename)

@patch("requests.post")
@patch.dict(os.environ, {"JWT": "test_token"}, clear=True)
def test_write_to_db_with_invalid_filename(mock_post):
    # Test with invalid filename
    polygon = Polygon([(0, 0), (1, 1), (1, 0)])
    area = 0.5
    tile_date = "2022-01-01"
    filename = 123
    with pytest.raises(TypeError):
        write_to_db(polygon, area, tile_date, filename)

@patch("requests.post")
@patch.dict(os.environ, {"JWT": "test_token"}, clear=True)
def test_write_to_db_with_request_exception(mock_post):
    # Test with a request exception
    mock_post.return_value.raise_for_status.side_effect = (
        requests.exceptions.RequestException
    )
    polygon = Polygon([(0, 0), (1, 1), (1, 0)])
    area = 0.5
    tile_date = "2022-01-01"
    filename = "test.tif"
    with pytest.raises(requests.exceptions.RequestException):
        write_to_db(polygon, area, tile_date, filename)


import torch as th

def _get_first_singular_vectors_power_method(X, Y, max_iter=500, tol=1e-06):
    eps = th.finfo(th.float64).eps
    try:
        y_score = next(col for col in Y.T if th.any(th.abs(col) > eps))
    except StopIteration as e:
        raise StopIteration("Y residual is constant") from e

    x_weights_old = 100
    for i in range(max_iter):
        x_weights = th.matmul(X.T, y_score) / th.matmul(y_score, y_score)
        x_weights /= th.sqrt(th.matmul(x_weights, x_weights)) + eps
        x_score = th.matmul(X, x_weights)
        y_weights = th.matmul(Y.T, x_score) / th.matmul(x_score.T, x_score)
        y_score = th.matmul(Y, y_weights) / (th.matmul(y_weights, y_weights) + eps)
        x_weights_diff = x_weights - x_weights_old
        if th.matmul(x_weights_diff, x_weights_diff) < tol or Y.shape[1] == 1:
            break
        x_weights_old = x_weights
    n_iter = i + 1
    return x_weights, y_weights, n_iter

def _get_first_singular_vectors_svd(X, Y):
    C = th.matmul(X.T, Y)
    U, _, Vt = th.linalg.svd(C)
    return U[:, 0], Vt[0, :]

def _svd_flip_1d(u, v):
    biggest_abs_val_idx = th.argmax(th.abs(u))
    sign = th.sign(u[biggest_abs_val_idx])
    u *= sign
    v *= sign


import requests
import re
headers = {
    'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/69.0.3497.100 Safari/537.36'
}
def page(offset):
    url = "http://maoyan.com/board/4?offset="+offset
    response = requests.get(url,headers=headers)
    html = response.text
    infos = re.findall(r'
(.*?)
',html,re.S)
    file = open("猫眼电影.txt","a+",encoding="utf-8")
    for info in infos:
        xuhao = re.findall(r'(.*?)',info,re.S)[0].strip()
        img = re.findall(r'',info,re.S)[0].strip()
        name = re.findall(r'
(.*?)',info,re.S)[0].strip()
        star = re.findall(r'
(.*?)
',info,re.S)[0].strip()
        date = re.findall(r'
(.*?)
',info,re.S)[0].strip()
        score1 = re.findall(r'(.*?)',info,re.S)[0].strip()
        score2 = re.findall(r'(.*?)',info,re.S)[0].strip()
        score = score1+score2
        print(xuhao,img,name,star,date,score)
        di = {
            "xuhao":xuhao,
            "img":img,
            "movieName":name,
            "star":star[3:],
            "date":date[5:],
            "score":score
        }
        print(di)
        # file.write(xuhao+"\t"+img+"\t"+name+"\t"+star+"\t"+date+"\t"+score+"\n\n")
        file.write(str(di)+"\n\n")
    file.close()

for i in range(10):
    page(offset=str(i*10))


"""""""""""""""""""""""""""""""""""""""""""""""""""""""""""
HDLTex: Hierarchical Deep Learning for Text Classification
script to run main fucntion and create hierarchical structure

* Copyright (C) 2018  Kamran Kowsari <kk7nc@virginia.edu>
* Last Update: Oct 26, 2018
* This file is part of  HDLTex project, University of Virginia.
* Free to use, change, share and distribute source code of RMDL
* Refrenced paper : HDLTex: Hierarchical Deep Learning for Text Classification
* Link: https://doi.org/10.1109/ICMLA.2017.0-134
* Comments and Error: email: kk7nc@virginia.edu
"""""""""""""""""""""""""""""""""""""""""""""""""""""""""""

import os
os.environ['KERAS_BACKEND'] = 'tensorflow'

import numpy as np
from keras.models import Sequential
import Data_helper
import BuildModel

if __name__ == "__main__":
    MEMORY_MB_MAX = 1600000 # maximum memory you can use
    MAX_SEQUENCE_LENGTH = 500 # Maximum sequance lenght 500 words
    MAX_NB_WORDS = 55000 # Maximum number of unique words
    EMBEDDING_DIM = 100 #embedding dimension you can change it to {25, 100, 150, and 300} but need to change glove version
    batch_size_L1 = 64 # batch size in Level 1
    batch_size_L2 = 64 # batch size in Level 2
    epochs = 1

    L1_model =0 # 0 is DNN, 1 is CNN, and 2 is RNN for Level 1
    L2_model =2 # 0 is DNN, 1 is CNN, and 2 is RNN for Level 2

    np.set_printoptions(threshold=np.inf)
    '''
    location of input data in two ways
    1: Tokenizer that is using GLOVE
    1: loadData that is using couting words or tf-idf
    '''

    X_train, y_train, X_test, y_test, content_L2_Train, L2_Train, content_L2_Test, L2_Test, number_of_classes_L2,word_index, embeddings_index,number_of_classes_L1 =  \
            Data_helper.loadData_Tokenizer(MAX_NB_WORDS,MAX_SEQUENCE_LENGTH)

    X_train_DNN, y_train_DNN, X_test_DNN, y_test_DNN, content_L2_Train_DNN, L2_Train_DNN, content_L2_Test_DNN, L2_Test_DNN, number_of_classes_L2_DNN = Data_helper.loadData()
    print("Loading Data is Done")
    #######################DNN Level 1########################
    if L1_model == 0:
        print('Create model of DNN')
        model = BuildModel.buildModel_DNN(X_train_DNN.shape[1], number_of_classes_L1, 8, 64, dropout=0.25)
        model.fit(X_train_DNN, y_train_DNN[:, 0],
                  validation_data=(X_test_DNN, y_test_DNN[:, 0]),
                  epochs=epochs,
                  verbose=2,
                  batch_size=batch_size_L1)

    #######################CNN Level 1########################
    if L1_model == 1:
        print('Create model of CNN')
        model = BuildModel.buildModel_CNN(word_index, embeddings_index,number_of_classes_L1,MAX_SEQUENCE_LENGTH,EMBEDDING_DIM,1)
        model.fit(X_train, y_train[:,0],
                  validation_data=(X_test, y_test[:,0]),
                  epochs=epochs,
                  verbose=2,
                  batch_size=batch_size_L1)
    #######################RNN Level 1########################
    if L1_model == 2:
        print('Create model of RNN')
        model = BuildModel.buildModel_RNN(word_index, embeddings_index,number_of_classes_L1,MAX_SEQUENCE_LENGTH,EMBEDDING_DIM)
        model.fit(X_train, y_train[:,0],
                  validation_data=(X_test, y_test[:,0]),
                  epochs=epochs,
                  verbose=2,
                  batch_size=batch_size_L1)

    HDLTex = [] # Level 2 models is list of Deep Structure
    ######################DNN Level 2################################
    if L2_model == 0:
        for i in range(0, number_of_classes_L1):
            print('Create Sub model of ',i)
            HDLTex.append(Sequential())
            HDLTex[i] = BuildModel.buildModel_DNN(content_L2_Train_DNN[i].shape[1], number_of_classes_L2_DNN[i],2, 1024, dropout=0.5)
            HDLTex[i].fit(content_L2_Train_DNN[i], L2_Train_DNN[i],
                      validation_data=(content_L2_Test_DNN[i], L2_Test_DNN[i]),
                      epochs=epochs,
                      verbose=2,
                      batch_size=batch_size_L2)
    ######################CNN Level 2################################
    if L2_model == 1:

        for i in range(0, number_of_classes_L1):
            print('Create Sub model of ', i)
            HDLTex.append(Sequential())
            HDLTex[i] = BuildModel.buildModel_CNN(word_index, embeddings_index,number_of_classes_L2[i],MAX_SEQUENCE_LENGTH,EMBEDDING_DIM,1)
            HDLTex[i].fit(content_L2_Train[i], L2_Train[i],
                          validation_data=(content_L2_Test[i], L2_Test[i]),
                          epochs=epochs,
                          verbose=2,
                          batch_size=batch_size_L2)
    ######################RNN Level 2################################
    if L2_model == 2:
        for i in range(0, number_of_classes_L1):
            print('Create Sub model of ', i)
            HDLTex.append(Sequential())
            HDLTex[i] = BuildModel.buildModel_RNN(word_index, embeddings_index,number_of_classes_L2[i],MAX_SEQUENCE_LENGTH,EMBEDDING_DIM)
            HDLTex[i].fit(content_L2_Train[i], L2_Train[i],
                          validation_data=(content_L2_Test[i], L2_Test[i]),
                          epochs=epochs,
                          verbose=2,
                          batch_size=batch_size_L2)


from os import listdir
from UI.FragmentDictionaryUI import Ui_Form
from Dialogs.InfoDialog import InfoDialog
from Dialogs.InputDialog import InputDialog
from Dialogs.ConfirmationDialog import ConfirmationDialog
from db_helper import Helper
from Elements.DictionaryTitle import DictionaryTitle
from Elements.DictionaryPair import DictionaryPair
from PyQt5 import QtCore, QtWidgets
import threading
import pyttsx3

# Словарь пользователя
class DictFragment(QtWidgets.QWidget, Ui_Form):
    def __init__(self, stacked_widget_top_bar: QtWidgets.QStackedWidget, top_bar_widget1,
                 top_bar_widget2):
        super().__init__()
        self.setupUi(self)
        self.engine = pyttsx3.init()
        voice_id = "HKEY_LOCAL_MACHINE\SOFTWARE\Microsoft\Speech\Voices\Tokens\TTS_MS_EN-US_ZIRA_11.0"
        self.engine.setProperty('voice', voice_id)

        # Верхняя панель
        self.stacked_widget_top_bar = stacked_widget_top_bar

        # Элемент верхней панели в окне с названиями тем словарей
        self.top_bar_widget1 = top_bar_widget1
        # Кнопки этой панели
        self.top_bar_widget1.pushButtonAddTheme.clicked.connect(self.button_add_theme_clicked)
        self.top_bar_widget1.pushButtonDeleteTheme.clicked.connect(self.button_delete_theme_clicked)

        # Элемент верхней панели в окне со словами словаря
        self.top_bar_widget2 = top_bar_widget2
        # Кнопки этой панели
        self.top_bar_widget2.pushButtonGoBack.clicked.connect(self.button_back_clicked)
        self.top_bar_widget2.pushButtonImport.clicked.connect(self.import_file)
        self.top_bar_widget2.pushButtonExport.clicked.connect(self.export_file)

        # Элемент предназначенный для ввода слов в словарь
        self.wordInput.pushButtonAddToDictionary.clicked.connect(self.button_add_word_clicked)
        self.wordInputAnimation = QtCore.QPropertyAnimation(self.wordInput, b'maximumHeight')
        self.wordInputAnimation.setEasingCurve(QtCore.QEasingCurve.OutBack)

        # Список тем
        self.listWidget_titles.itemSelectionChanged.connect(self.item_selected)
        self.listWidget_titles.setUniformItemSizes(True)
        # Список слов
        self.listWidget_words.setSpacing(2)
        self.listWidget_words.setUniformItemSizes(True)

        self.selected_item_index = None  # Индекс выбранного элемента в списке тем
        self.top_panel_page = 1  # Страница верхней панели
        self.current_title = None  # Выбранная тема. Используется для отображения слов во втором окне
        self.button_to_label = dict()  # Соответствие кнопки к заголовку темы
        self.button_to_index = dict()  # Порядковый номер кнопки удаления слов
        self.button_speech_to_word = dict()  # Слово, соответствующее кнопке озвучивания
        self.db_helper = Helper()  # Помощник по бд

        self.theme_titles = []  # Список заголовков тем
        self.words = []  # Список слов определенной тем
        self.translations = []  # Список переводов слов определенной темы

        self.list_widget1_set_data()  # Заполнение списка тем

    # Добавление новой темы
    def button_add_theme_clicked(self):
        # Вызов диалогового окна ввода и получение результата
        dialog = InputDialog()
        value = dialog.get_text()
        if value is not None and len(value) != 0:  # Проверка входных данных
            # Если тема уже содержится в базе данных, то к названию добавляется цифра
            if value in self.theme_titles:
                c = 2
                while True:
                    if value + '_' + str(c) not in self.theme_titles:
                        value = value + '_' + str(c)
                        break
                    c += 1
            # Добавление в базу данных и обновление списка
            self.db_helper.add_title(value)
            self.list_widget1_set_data()

    # Заполнение списка тем
    def list_widget1_set_data(self):
        self.listWidget_titles.clear()
        self.button_to_label.clear()
        # Обновление списка тем
        self.theme_titles = self.db_helper.get_titles()
        # Заполнение списка элементами
        for i in range(len(self.theme_titles)):
            item = QtWidgets.QListWidgetItem()
            item.setData(0, i)
            wdt = DictionaryTitle()
            wdt.pushButton.setText(f'Перейти')
            wdt.pushButton.clicked.connect(self.button_goto_theme_clicked)
            wdt.label.setText(str(self.theme_titles[i]))
            self.button_to_label[wdt.pushButton] = str(self.theme_titles[i])
            item.setSizeHint(wdt.sizeHint())
            self.listWidget_titles.addItem(item)
            self.listWidget_titles.setItemWidget(item, wdt)

    # Удаление темы
    def button_delete_theme_clicked(self):
        # Проверка на наличие выбранного элемента
        if self.selected_item_index is not None:
            # Подтверждение удаления
            confirmation = ConfirmationDialog(
                "Действительно удалить?").get_results()  # Вызов диалогового окна
            if not confirmation:
                return
            # Удаление из базы данных и обновление списка
            self.db_helper.delete_title(self.theme_titles[self.selected_item_index])
            self.list_widget1_set_data()
            self.selected_item_index = None

    # Переход к выбранной теме
    def button_goto_theme_clicked(self):
        self.stackedWidget.setCurrentIndex(1)  # Переключение страницы основного окна
        self.stacked_widget_top_bar.setCurrentIndex(2)  # Переключение страницы верхней панели
        self.word_input_animation(True)  # Анимация поля ввода слов
        self.top_panel_page = 2
        self.current_title = (self.button_to_label[self.sender()])
        # Установка элементов в список
        self.list_widget2_set_data_words()

    # Переход со словаря к списку тем
    def button_back_clicked(self):
        self.stackedWidget.setCurrentIndex(0)
        self.stacked_widget_top_bar.setCurrentIndex(1)
        self.word_input_animation(False)  # Обратная анимация поля ввода слов
        self.top_panel_page = 1
        self.current_title = None

    # Добавление слова в словарь
    def button_add_word_clicked(self):
        # Считывание данных с полей ввода
        word = self.wordInput.lineEditWord.text().strip()
        translation = self.wordInput.lineEditTranslation.text().strip()
        # Проверка введенных данных
        if len(word.strip()) != 0 and len(translation) != 0:
            # Проверка на содержание в базе данных
            if word not in self.words and translation not in self.translations:
                # Добавление в базу данных
                self.db_helper.add_word(self.current_title, word, translation)
                # Очистка полей ввода
                self.wordInput.lineEditWord.setText('')
                self.wordInput.lineEditTranslation.setText('')
                # Обновление списка и прокручивание его вниз
                self.list_widget2_set_data_words()
                self.listWidget_words.scrollToBottom()

    # Заполнение словаря
    def list_widget2_set_data_words(self):
        # Обновление списков
        self.words, self.translations = self.db_helper.get_words_by_title(self.current_title)
        self.listWidget_words.clear()
        for i in range(len(self.words)):
            item = QtWidgets.QListWidgetItem()
            item.setData(0, i + 1)
            wdt = DictionaryPair()
            wdt.label.setText(str(i + 1))
            wdt.label_2.setText(self.words[i])
            wdt.label_3.setText(self.translations[i])
            item.setSizeHint(wdt.sizeHint())
            self.button_to_index[wdt.pushButton] = i
            self.button_speech_to_word[wdt.pushButtonSpeech] = self.words[i]
            wdt.pushButton.clicked.connect(self.button_delete_word_clicked)
            wdt.pushButtonSpeech.clicked.connect(self.voice)
            self.listWidget_words.addItem(item)
            self.listWidget_words.setItemWidget(item, wdt)

    # Удаление слова из словаря
    def button_delete_word_clicked(self):
        index = self.button_to_index[self.sender()]  # Номер нажатой кнопки
        # Удаление из базы данных и обновление списка
        self.db_helper.delete_word(self.current_title, self.words[index], self.translations[index])
        self.list_widget2_set_data_words()

    # Сохранение номера выбранного элемента списка тем
    def item_selected(self):
        if len(self.listWidget_titles.selectedItems()):
            self.selected_item_index = self.listWidget_titles.selectedItems()[0].data(0)

    # Анимация поля ввода слов
    def word_input_animation(self, direction):
        self.wordInputAnimation.setStartValue(0 if direction else 55)
        self.wordInputAnimation.setEndValue(55 if direction else 0)
        self.wordInputAnimation.setDuration(250)
        self.wordInputAnimation.start()

    def voice(self):
        threading.Thread(
            target=self.text_to_speech, args=(self.sender(),), daemon=True
        ).start()

    # Озвучивание слова
    def text_to_speech(self, sender):
        self.engine.say(self.button_speech_to_word[sender])
        try:
            self.engine.runAndWait()
        except RuntimeError:
            pass

    # Ввод данных в словарь из
    def import_file(self):
        # Выбор файла
        filename = QtWidgets.QFileDialog.getOpenFileName(
            self, 'Выбрать файл', '', '(*.txt);; Все файлы (*)')[0]
        # Проверка названия файла на пустоту
        if len(filename) == 0:
            return
        file = open(filename, encoding='utf8')
        data = file.read()  # Чтение данных из файла
        file.close()

        # Проверка содержимого файла на пустоту
        if len(data.splitlines()) == 0:
            info_dialog = InfoDialog('Выбранный файл пуст. Изменения не внесены')
            info_dialog.exec_()
            return

        f_words, f_translations = [], []  # Данные из файла
        c = 1
        # Построчное чтение файла
        for line in data.splitlines():
            try:
                if '"' in line:
                    raise IncorrectSymbol
                f_word, f_translation = line.split(';')  # Чтение слова и перевода слова из строки
            except ValueError:  # В случае ошибки вызывается диалоговое окно с информацией
                info_dialog = InfoDialog(f'Ошибка в строке {c}. Изменения не внесены')
                info_dialog.exec_()
                return
            except IncorrectSymbol:
                info_dialog = InfoDialog(
                    f'Недопустимый символ (") в строке {c}. Изменения не внесены')
                info_dialog.exec_()
                return
            except Exception:
                info_dialog = InfoDialog(f'Ошибква в файле')
                info_dialog.exec_()
                return
            f_words.append(f_word)
            f_translations.append(f_translation)
            c += 1

        # Подтверждение
        confirmation_dialog = ConfirmationDialog('Данные словаря будут переписаны.\nПродолжить?')
        result = confirmation_dialog.get_results()

        # Выход из функции при отклонении подтверждения
        if not result:
            return

        # Удаление старых слов из базы данных
        for word, translation in zip(self.words, self.translations):
            self.db_helper.delete_word(self.current_title, word, translation)
        # Внесение новых слов в базу данных
        for word, translation in zip(f_words, f_translations):
            self.db_helper.add_word(self.current_title, word, translation)
        # Обновление списка
        self.list_widget2_set_data_words()

    # Сохранение словаря в файл
    def export_file(self):
        # Выбор директории
        pathname = QtWidgets.QFileDialog.getExistingDirectory()
        # Выход из функции если директории нет
        if len(pathname) == 0:
            return

        files = listdir(pathname)  # Список всех файлов в выбранной директории
        filename = self.current_title  # Название сохраняемого файла
        # Проверка на наличие файла с таким же названием в выбранной директории
        if filename + '.txt' in files:
            c = 2
            # Добавление числа к файлу, до тех пор, пока название не окажется свободным
            while True:
                if filename + '_' + str(c) + '.txt' not in files:
                    break
                c += 1
            filename = filename + '_' + str(c)
        filename += '.txt'
        file = open(filename, 'a', encoding='utf8')
        for word, translation in zip(self.words, self.translations):
            file.write(
                word + ";" + translation + "\n")  # Построчная запись в файл в виде слово;перевод
        file.close()

class IncorrectSymbol(Exception):
    pass


import mlflow
from prefect import flow, get_run_logger, task
from hyperopt import fmin, tpe, STATUS_OK, Trials, space_eval
from prefect.task_runners import SequentialTaskRunner
import pandas as pd
from typing import Dict
import xgboost
from churnlib import create_xgboost_churn_model, xgb_space, CHURN_MLFLOW_EXPERIMENT_NAME
from sklearn.metrics import (
    roc_auc_score,
    accuracy_score,
    precision_score,
    recall_score,
    f1_score,
)

TRAINING_DATA_PATH = "../data/Bank Customer Churn Prediction.csv"
mlflow.set_tracking_uri("http://localhost:5000")
mlflow.set_experiment(CHURN_MLFLOW_EXPERIMENT_NAME)

@task
def ingest_data():
    """Ingest data from a source"""
    logger = get_run_logger()
    logger.info("Ingesting data...")
    churn_db = pd.read_csv(TRAINING_DATA_PATH)
    return churn_db

@task
def optimize_hyperparameters(data: pd.DataFrame, search_space: Dict):
    """Optimize hyperparameters for a model"""
    logger = get_run_logger()
    logger.info("Optimizing hyperparameters...")

    def xgboost_objective_function(params: Dict) -> float:

        xgb_model = create_xgboost_churn_model(
            data=data,
            params={},
            target_col="churn",
            exclude_cols=["customer_id", "country", "gender"],
        )
        train = xgboost.DMatrix(
            xgb_model.split_data["x_train"], xgb_model.split_data["y_train"]
        )
        res = xgboost.cv(
            params,
            train,
            num_boost_round=100,
            nfold=5,
            metrics={"auc"},
            seed=0,
            callbacks=[
                xgboost.callback.EvaluationMonitor(show_stdv=True),
                xgboost.callback.EarlyStopping(15),
            ],
        )
        best_loss = res["test-auc-mean"].iloc[-1]
        return {"loss": best_loss, "status": STATUS_OK}

    trials = Trials()
    best = fmin(
        fn=xgboost_objective_function,
        space=search_space,
        algo=tpe.suggest,
        max_evals=3,
        trials=trials,
    )
    return space_eval(search_space, best)

@flow(task_runner=SequentialTaskRunner())
def churn_model_training():
    """Train a model with the given parameters"""
    logger = get_run_logger()
    logger.info("Training model...")
    churn_db = ingest_data()
    best_hyperparameters = optimize_hyperparameters(data=churn_db, search_space=xgb_space)

    mlflow.xgboost.autolog()
    with mlflow.start_run():

        churn_model = create_xgboost_churn_model(
            data=churn_db,
            params={},
            target_col="churn",
            exclude_cols=["customer_id", "country", "gender"],
        )
        x_train, y_train = (
            churn_model.split_data["x_train"],
            churn_model.split_data["y_train"],
        )
        x_test, y_test = (
            churn_model.split_data["x_test"],
            churn_model.split_data["y_test"],
        )
        best_xgb = xgboost.XGBClassifier(**best_hyperparameters)
        best_xgb.fit(x_train, y_train)

        predictions = best_xgb.predict(x_test)

        accuracy_metric = accuracy_score(y_test, predictions)
        precision_metric = precision_score(y_test, predictions)
        recall_metric = recall_score(y_test, predictions)
        f1_score_metric = f1_score(y_test, predictions)

        mlflow.log_metric(
            "AUC score", roc_auc_score(y_test, best_xgb.predict_proba(x_test)[:, 1])
        )

        mlflow.end_run()

    return None

if __name__ == "__main__":
    churn_model_training()


from django.db import models

from wagtail.core.models import Page
from wagtail.core.fields import StreamField
from wagtail.admin.edit_handlers import (
  MultiFieldPanel,
  StreamFieldPanel,
  FieldPanel,
  InlinePanel,
)
from wagtail.images.edit_handlers import ImageChooserPanel
from modelcluster.fields import ParentalKey

from ..blocks.questions import QuestionsPageStreamBlock
from ...variables import QUESTION_CATEGORY_CHOICES

class QuestionsOverviewCategory(models.Model):

  page = ParentalKey('home.QuestionsOverview', on_delete=models.CASCADE, related_name='categories')

  name = models.CharField(max_length=164)
  description = models.TextField()
  hide = models.BooleanField(default=True)

QuestionsOverviewCategory.panels = [
  FieldPanel('name', classname='col12'),
  FieldPanel('description', classname='col12'),
  FieldPanel('hide', classname='col6')
]

class QuestionsOverview(Page):

  template = 'home/questions/questions_overview.html'

  cover_image = models.ForeignKey(
    'wagtailimages.Image',
    on_delete=models.SET_NULL,
    related_name='+',
    null=True,
    blank=False
  )
  info_text = models.TextField(verbose_name='description', null=True)

  # last
  def featured(self):
    return QuestionsPage.objects.live() \
      .order_by('-first_published_at')[0]

  # All but last
  def others(self):

    count = QuestionsPage.objects.live().count()

    if count > 1:
      return QuestionsPage.objects.live() \
        .order_by('-first_published_at')[1:count]

    return None

  def get_context(self, request):
    context = super().get_context(request)
    context['previous'] = request.META.get('HTTP_REFERER')
    return context

QuestionsOverview.content_panels = Page.content_panels + [
  FieldPanel('info_text'),
  ImageChooserPanel('cover_image'),
  MultiFieldPanel([
    InlinePanel('categories')
  ], heading='Categories')
]

QuestionsOverview.parent_page_types = ['home.HomePage']
QuestionsOverview.subpage_types = ['home.QuestionsPage']

class QuestionsPage(Page):

  template = 'home/questions/questions_page.html'

  cover_image = models.ForeignKey(
    'wagtailimages.Image',
    on_delete=models.SET_NULL,
    related_name='+',
    null=True,
    blank=True
  )
  category = models.IntegerField(choices=QUESTION_CATEGORY_CHOICES, default=1)
  content = StreamField(QuestionsPageStreamBlock(), null=True)

QuestionsPage.content_panels = Page.content_panels + [
  MultiFieldPanel([
    ImageChooserPanel('cover_image', classname='col12'),
    FieldPanel('category', classname='col12')
  ], heading='Info'),
  MultiFieldPanel([
    StreamFieldPanel('content')
  ], heading='Content')
]

QuestionsPage.parent_page_types = ['home.QuestionsOverview']
QuestionsPage.subpage_types = []


"""
Author: Wouter Van Gansbeke, Simon Vandenhende
Licensed under the CC BY-NC 4.0 license (https://creativecommons.org/licenses/by-nc/4.0/)
"""
import numpy as np
import torch
from torch.utils.data import Dataset
import os
import json
import random
from PIL import Image
import pdb

def unpickle(file):
    import _pickle as cPickle
    with open(file, 'rb') as fo:
        dict = cPickle.load(fo, encoding='latin1')
    return dict

"""
    AugmentedDataset
    Returns an image together with an augmentation.
"""
class AugmentedDataset(Dataset):
    def __init__(self, dataset):
        super(AugmentedDataset, self).__init__()
        transform = dataset.transform
        dataset.transform = None
        self.dataset = dataset

        if isinstance(transform, dict):
            self.image_transform = transform['standard']
            self.augmentation_transform = transform['augment']

        else:
            self.image_transform = transform
            self.augmentation_transform = transform

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

    def __getitem__(self, index):
        sample = self.dataset.__getitem__(index)
        image = sample['image']
        sample['image'] = self.image_transform(image)
        sample['image_augmented'] = self.augmentation_transform(image)
        return sample

"""
    NeighborsDataset
    Returns an image with one of its neighbors.
"""
class NeighborsDataset(Dataset):
    def __init__(self, dataset, indices, num_neighbors=None, predicted_labels=None):
        super(NeighborsDataset, self).__init__()
        transform = dataset.transform

        if isinstance(transform, dict):
            self.anchor_transform = transform['standard']
            self.neighbor_transform = transform['augment']
        else:
            self.anchor_transform = transform
            self.neighbor_transform = transform

        dataset.transform = None
        self.dataset = dataset
        self.indices = indices # Nearest neighbor indices (np.array  [len(dataset) x k])
        self.predicted_labels = predicted_labels
        if num_neighbors is not None:
            self.indices = self.indices[:, :num_neighbors+1]

        assert(self.indices.shape[0] == len(self.dataset))

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

    def __getitem__(self, index):
        output = {}
        anchor = self.dataset.__getitem__(index)

        if self.predicted_labels is not None:
            label = self.predicted_labels[index]
            neighbor_indices = self.indices[index]
            neighbor_labels = self.predicted_labels[neighbor_indices]
            same_label_neighbors = np.where(neighbor_labels == label)[0]
            valid_neighbors = neighbor_indices[same_label_neighbors]
            if valid_neighbors.size == 0:
                valid_neighbors = neighbor_indices
        else:
            valid_neighbors = self.indices[index]
        neighbor_index = np.random.choice(valid_neighbors, 1)[0]
        neighbor = self.dataset.__getitem__(neighbor_index)

        anchor['image'] = self.anchor_transform(anchor['image'])
        neighbor['image'] = self.neighbor_transform(neighbor['image'])

        output['anchor'] = anchor['image']
        output['neighbor'] = neighbor['image']
        output['possible_neighbors'] = torch.from_numpy(self.indices[index])
        output['target'] = anchor['target']

        return output

"""
    NoisyDataset
"""
class NoisyDataset(Dataset):
    def __init__(self, dataset, meta_info):
        super(NoisyDataset, self).__init__()
        self.transform = dataset.transform
        dataset.transform = None
        self.dataset = dataset
        self.probability = meta_info['probability']
        self.mode = meta_info['mode']

        r = meta_info['r']
        pred = meta_info['pred']
        noise_file = meta_info['noise_file']
        noise_mode = meta_info['noise_mode']
        dataset_name = meta_info['dataset']

        if os.path.exists(noise_file):
            noise = json.load(open(noise_file,"r"))
            noise_labels = noise['noise_labels']
        else:
            import pdb;pdb.set_trace()
            raise NotImplementedError()

        self.noise_labels = noise_labels
        if self.mode == 'labeled':
            self.indices = pred.nonzero()[0]
        elif self.mode == 'unlabeled':
            self.indices = (1-pred).nonzero()[0]
        elif self.mode == 'all' or self.mode == 'pretext':
            self.indices = list(range(len(self.dataset)))
        elif self.mode == 'neighbor':
            # self.indices = pred.nonzero()[0]
            self.indices = list(range(len(self.dataset)))
        elif self.mode == 'proportional':
            temp_indices = list(range(len(self.dataset)))
            self.indices = [i for i in temp_indices if i not in meta_info['idx_remove']]
        else:
            raise ValueError('Invalid noisy dataset mode')

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

    def __getitem__(self, index):
        true_index = self.indices[index]
        image = self.dataset.__getitem__(true_index)
        if self.mode=='labeled':
            img, target, prob = image, self.noise_labels[true_index], self.probability[true_index]
            img = img['image']
            img1 = self.transform(img)
            img2 = self.transform(img)
            return img1, img2, target, prob
        elif self.mode=='unlabeled':
            img = image
            img = img['image']
            img1 = self.transform(img)
            img2 = self.transform(img)
            return img1, img2
        elif self.mode=='all':
            img, target = image, self.noise_labels[true_index]
            img = img['image']
            img = self.transform(img)
            return img, target, true_index
        elif self.mode=='neighbor':
            return image
        elif self.mode=='pretext':
            img = image['image']
            if self.transform is not None:
                img = self.transform(img)
            return {'image': img, 'target': self.noise_labels[true_index]}
        elif self.mode=='proportional':
            img, target, prob = image, self.noise_labels[true_index], self.probability[true_index]
            img = img['image']
            img1 = self.transform(img)
            img2 = self.transform(img)
            return img1, img2, target, prob


# Import libaries
from dataclasses import dataclass, field
from typing import List
import csv

#/////////////////////////////////////////////////////////////////#

# Equivalent to c structs or OOP classes
@dataclass
class Product:
    name: str
    price: float = 0.0

@dataclass
class ProductStock:
    product: Product
    quantity: int

@dataclass
class Shop:
    cash: float = 0.0
    stock: List[ProductStock] = field(default_factory=list)

@dataclass
class Customer:
    name: str = ""
    budget: float = 0.0
    shoppingList: List[ProductStock] = field(default_factory=list)

#/////////////////////////////////////////////////////////////////#

# Function to stock the shop by reading in the stock.csv file
def createAndStockShop():
    s = Shop()

    # Try to open and read the csv file
    try:
        with open("../stock.csv") as csv_file:
            csv_reader = csv.reader(csv_file, delimiter=',')
            first_row = next(csv_reader)
            s.cash = float(first_row[0]) # First row = shop cash
            # Iterate through rows and create Product, ProductStock and stock instances
            for row in csv_reader:
                p = Product(row[0], float(row[1]))
                ps = ProductStock(p, int(row[2]))
                s.stock.append(ps)
        return s # Return shop datacase

    except FileNotFoundError:
        print(f"Cannot open {csv_file}")

# Function to print product name and price
def printProduct(p):
    print(f'{p.name}\t\t{p.price:,.2f}')

# Function to display the shop
def printShop(s):
    print('##########################################################\n')
    print('\t\tGENERAL STORE OVERVIEW\t\t\n')
    print('##########################################################\n')
    print(f"Shop Cash: {s.cash:.2f}")
    print("*********************************************\n")
    print("\t\tCatalogue")
    print("*********************************************")
    print("\n Product\t\tPrice\t\tQty")
    print("_____________________________________________\n")

    #  Loop through each stock item and print name, price and qty
    for item in s.stock:
        print(f"{item.product.name} \t\t{item.product.price:.2f}\t\t{item.quantity}\n")
        print("----------------------------------------------\n")

# Function to read a customer's order from a csv file
def customerOrder():

    filepath =  input("Enter csv filepath: ") # Get csv filepath

    # Try to open and read csv file
    try:
        with open(filepath) as csv_file:
            csv_reader = csv.reader(csv_file, delimiter=',')
            first_row = next(csv_reader)
            c = Customer(first_row[0], float(first_row[1])) # Create customer class
            # Create customer shoppingList by iterating through the rest of the csv
            for row in csv_reader:
                name = row[0]
                quantity = int(row[1])
                p = Product(name)
                ps = ProductStock(p, quantity)
                c.shoppingList.append(ps) # Add to shoppingList

            return c # Return customer

    except FileNotFoundError:
        print(f"Cannot open {filepath}")

# Function to print customer
def printCustomer(c,s):

    total = 0 # Variable to track customer order total

    print("##########################################################\n")
    print(f" Customer: {c.name}\n Budget:   {c.budget:.2f}\n")
    print("***************************************************\n")
    print("\n Product\t\tPrice\tQty\tUnit Total")
    print("___________________________________________________\n")

    # Loop through shoppingList
    for item in c.shoppingList:

        getPrice(c, s) # Calling the getPrice function
        cost = 0 # Variable to track unit cost of a product

        # Comparing customer items to see if it exists in the shop stock
        for customerItem in s.stock:
            if item.product.name.lower() == customerItem.product.name.lower():
                # Get the unit cost of the product (quantity x price) if it's in stock. Cost is 0 if item is not in stock
                cost =+ item.quantity * customerItem.product.price

        print(f"{item.product.name}\t\t{item.product.price:.2f}\t{item.quantity}\t{cost:.2f}")
        print("---------------------------------------------------\n")

        total += cost # Update total cost

    print(f"Total\t\t\t\t\t{total:.2f}")
    print("---------------------------------------------------\n")

# Function to get a product price
def getPrice(c, s):
    # Loops through the shop stock to find a match and retrieves the product price
    for shopProduct in s.stock:
        for item in c.shoppingList:
            if (item.product.name.lower() == shopProduct.product.name.lower()):
                item.product.price = shopProduct.product.price

# Function to process customer order
def processOrder(shop, customer):

    # Variable to track whether to checkout a customers order
    checkout = input(f"\nCheckout {customer.name}'s order? (y/n): ").lower()

    if checkout == 'y':

        print("\n -------------------------------------------------\n")
        print("\n Checkout complete!\n")

        # Iterate over the customer shoppingList
        for item in customer.shoppingList:

            # Process order
            for customerItem in s.stock:
                # if the item is a shop stock item calculate order total
                if item.product.name.lower() == customerItem.product.name.lower():

                    orderTotal = item.quantity * customerItem.product.price

                    # If customer orders less than or equal to the the shop has in stock
                    if item.quantity <= customerItem.quantity:
                        # and the order total is less than or equal the customers budget
                        if customer.budget > orderTotal:
                            # Update the shop stock
                            customerItem.quantity -= item.quantity
                            # Update the shop and customers cash
                            shop.cash += orderTotal
                            customer.budget -= orderTotal

                        else: # If the order total is more than the customers budget
                            print("\n***************************************************************\n")
                            print("\nCannot complete order: insufficient Funds")
                            print("\n***************************************************************\n")

                    else: # If customer orders more than is stocked in the shop
                        print("\n***************************************************************\n");
                        print("\n !!! Order incomplete: due to {} stock shortage !!!\n".format(item.product.name));
                        print("\n***************************************************************\n");

#  Function to get customer name, budget and shopping list (live order)
def liveCustomer(shop):

    # Get customer name ad budget from input
    customerName = input("Welcome to the general store. What's your name? ")
    customerBudget = float(input(f"Hi {customerName}, what's your budget? "))

    print("\nPerfect. Please choose from our product list:\n")
    print("*********************************************\n")
    print("\t\tCatalogue")
    print("_____________________________________________")
    print("\n Product\t\tPrice\n")
    print("*********************************************\n")

    # Loop through shop stock and print product name and price
    for i in shop.stock:
        printProduct(i.product)
        print("---------------------------------------------\n")

    live_customer = Customer(customerName, customerBudget) # Create Customer class instance

    addItem = 'y' # Varaible to track if customer wishes to add additional products

    # Continue the while loop until the customer is finished adding items by pressing 'n' to adding additional items
    while addItem != "n":

    # Create customer shoppingList #

        # Get product name and order quantity from input
        customerItem = input("\nWhat product would you like? ").lower()

        while True:
            try:
                customerQuantity = int(input("How many? "))
                break
            except ValueError:
                print("Please enter your budget as an number\t")

        p = Product(customerItem) # Check product name
        ps = ProductStock(p, customerQuantity) # Check order quantity
        live_customer.shoppingList.append(ps) # Add to shoppingList

        # Continue the while loop until the customer is finished adding items
        addItem = input("Do you need anything else? (y/n): ").lower()

    return live_customer # Return customer

# Function to show main shop menu options
def menuOptions():

    print('##########################################################\n')
    print('\t\t\tGENERAL STORE\t\t\n')
    print('##########################################################\n')
    print("\nWelcome to the general store. How can we help you today?\n")
    print("\t-[1] View shop")
    print("\t-[2] Place an order")
    print("\t-[3] Import Order from a file")
    print("\t-[0] Exit")

# Function for main shop menu
def displayMenu(shop):

    menuOptions() # Show main menu

    while True:

        choice = input("\nPlease select an option: ") # Variable to track customer choice

        if choice == '1': # Display shop #

            printShop(shop) # Print the shop
            displayMenu(shop) # Return to main menu

        elif choice == '2': # Place a live order #

            live_customer = liveCustomer(shop) # Create a customer from live input
            printCustomer(live_customer, shop) # Print the customer
            processOrder(shop, live_customer ) # Process the order by calling processOrder
            print("\nThe balance on {}'s account:€{:.2f}\n".format(live_customer.name, live_customer.budget)) # Print customers new balance
            displayMenu(shop) # Return to main menu

        elif choice == '3': # Read order from a csv file #

            customer = customerOrder() # Create a customer from csv file
            printCustomer(customer, shop) # Print the customer
            processOrder(shop, customer) # Process the order by calling processOrder
            print("\nThe balance on {}'s account: EUR {:.2f}\n".format(customer.name, customer.budget)) # Print customers new balance
            displayMenu(shop) # Return to main menu

        elif choice == '0': # Exit programme #
            print("Thanks for shopping at the general store, goodbye!")
            exit()  # Break out of display menu loop

        else:
            displayMenu(shop) # Any other key is pressed show main menu

s = createAndStockShop() # Stock the shop
displayMenu(s) # Show the display menu