import copy
import sys
sys.path.extend([
    "./"
])
import numpy as np
import os
import optuna
import seaborn as sns
import hashlib
import uuid
import re
import argparse
import copy
from tqdm import tqdm
import pandas as pd
import datetime

from src.custom_metrics.features import FeatureStats
from src.custom_metrics.features import load_stats

from src.general_utils import util_path
from src.general_utils import util_general
from src.general_utils import util_ensemble
from src.general_utils import util_metric

# For aspect ratio 4:3.
sns.set_context("paper")
sns.set_theme(style="ticks")

def get_parser():

    parser = argparse.ArgumentParser(description='Ensemble search.')

    parser.add_argument('--source_dir', type=str, default='./reports/AIforCOVID', help='Directory name to fake samples.')
    parser.add_argument('--dataset_name', type=str, default='AIforCOVID',  choices=['retinamnist', 'pneumoniamnist', 'breastmnist', 'AIforCOVID', 'organamnist'], help='The name of dataset')
    parser.add_argument('--gpu_ids', type=str, default='0', help='gpu ids: e.g. 0  0,1,2, 0,2. use -1 for CPU')
    # GANs params.
    #parser.add_argument('--gan_models', type=util_general.parse_comma_separated_list, default='MHGAN,SNGAN,StyleGAN2-D2DCE,ReACGAN-ADA,ReACGAN-ADC,ReACGAN-DiffAug,ACGAN-Mod,ReACGAN,BigGAN-DiffAug,BigGAN-Info,StyleGAN2-DiffAug,ACGAN-Mod-TAC,BigGAN,ReACGAN-TAC,BigGAN-ADA,StyleGAN2-Info,ACGAN-Mod-ADC,StyleGAN2-ADA,ReACGAN-Info,StyleGAN2,ContraGAN,SAGAN', help='List of GANs to enable in the ensemble') #ACGAN-Mod-ADC,SAGAN,ReACGAN-ADA,StyleGAN2-DiffAug,StyleGAN2  # MHGAN,SNGAN,StyleGAN2-D2DCE,ReACGAN-ADA,ReACGAN-ADC,ReACGAN-DiffAug,ACGAN-Mod,ReACGAN,BigGAN-DiffAug,BigGAN-Info,StyleGAN2-DiffAug,ACGAN-Mod-TAC,BigGAN,ReACGAN-TAC,BigGAN-ADA,StyleGAN2-Info,ACGAN-Mod-ADC,StyleGAN2-ADA,ReACGAN-Info,StyleGAN2,ContraGAN,SAGAN
    parser.add_argument('--gan_models', type=util_general.parse_comma_separated_list, default='MHGAN,SNGAN,StyleGAN2-D2DCE,ReACGAN-ADA', help='List of GANs to enable in the ensemble') #ACGAN-Mod-ADC,SAGAN,ReACGAN-ADA,StyleGAN2-DiffAug,StyleGAN2  # MHGAN,SNGAN,StyleGAN2-D2DCE,ReACGAN-ADA,ReACGAN-ADC,ReACGAN-DiffAug,ACGAN-Mod,ReACGAN,BigGAN-DiffAug,BigGAN-Info,StyleGAN2-DiffAug,ACGAN-Mod-TAC,BigGAN,ReACGAN-TAC,BigGAN-ADA,StyleGAN2-Info,ACGAN-Mod-ADC,StyleGAN2-ADA,ReACGAN-Info,StyleGAN2,ContraGAN,SAGAN
    parser.add_argument('--gan_steps', type=util_general.parse_comma_separated_list, default='20000,40000',help='Iter or Iters to sample each GAN') #20000,40000 # 20000,40000,60000,80000,100000
    parser.add_argument("--post_resizer", type=str, default="friendly", help="which resizer will you use to evaluate GANs in ['legacy', 'clean', 'friendly']")
    parser.add_argument('--eval_backbone', type=str, default='ResNet50_torch',  choices=['InceptionV3_torch', 'ResNet50_torch', 'SwAV_torch', 'InceptionV3_torch__medical',  'InceptionV3_torch__truefake', 'ResNet50_torch__medical', 'ResNet50_torch__truefake'])
    parser.add_argument('--n_samples', type=int, default='664', help='Total number of images generated by the ensemble.')
    parser.add_argument('--split', type=str, default='train', choices=['train', 'val'])
    # Ensemble search params.
    parser.add_argument('--name_obj', type=str, default='dc_inter', choices=['dc_inter', 'pr_inter', 'fid_inter', 'dc_inter__dc_intra', 'pr_inter__pr_intra', 'fid_inter__fid_intra'])
    parser.add_argument('--pairwise_intra', type=bool, default=True)
    #parser.add_argument('--pairwise_intra_df', type=bool, default=True)
    parser.add_argument('--n_trial', type=int, default=10)
    parser.add_argument('--dim_reduction', type=bool, default=False)
    parser.add_argument('--synth_samples_reduction', type=int, default=100000)
    # Downstream task params.
    parser.add_argument("-cfg", "--cfg_file", type=str, default="./src/configs_downstream_task/pneumoniamnist.yaml")
    parser.add_argument('--init_w', type=str, default='uniform', choices=['uniform', 'random', 'fid'],  help='Weight initialization')
    parser.add_argument('--n_train', type=int, default=20, help='Number of times the experiment is repeated.')
    parser.add_argument("--num_workers", type=int, default=8)

    return parser

if __name__ == '__main__':

    parser = get_parser()
    args, unknown = parser.parse_known_args()

    # Directories.
    source_dir = args.source_dir
    dataset_name = args.dataset_name
    samples_dir = os.path.join(source_dir, 'samples')
    ensemble_dir = os.path.join(source_dir, 'ensemble')
    reports_dir = os.path.join(source_dir, 'downstream_task')
    cache_dir = os.path.join(source_dir, 'features')
    util_path.create_dir(reports_dir)

    # Parameters features.
    post_resizer = args.post_resizer
    eval_backbone = args.eval_backbone
    n_samples = args.n_samples
    split = args.split

    # Parameters ensemble.
    name_obj = args.name_obj
    n_obj = len(name_obj.split('__'))
    n_trial = args.n_trial
    pairwise_intra = args.pairwise_intra
    dim_reduction = args.dim_reduction
    synth_samples_reduction = args.synth_samples_reduction

    # Add to the filename the current time.
    filename = f'ensemble_search-n_trial__{n_trial}-name_obj__{name_obj}-pairwise_intra__{pairwise_intra}-dim_reduction__{dim_reduction}-synth_samples_reduction__{synth_samples_reduction}-{eval_backbone}-{post_resizer}-{n_samples}'
    now = datetime.datetime.now()
    #filename = f'{now.strftime("%Y-%m-%d_%H-%M-%S")}_{filename}'

    ensemble_dir = os.path.join(ensemble_dir, filename)
    util_path.create_dir(ensemble_dir)

    directions = []
    for i in name_obj.split('__'):
        if i == 'pr_inter':
            directions.append('maximize')
        elif i == 'pr_intra':
            directions.append('minimize')
        elif i == 'dc_inter':
            directions.append('maximize')
        elif i == 'dc_intra':
            directions.append('minimize')
        elif i=='fid_inter':
            directions.append('minimize')
        elif i=='fid_intra':
            directions.append('maximize')
        else:
            raise NotImplementedError

    # Load real features.
    data_real_list = []
    stats_real = load_stats(cache_dir=os.path.join(cache_dir, f'real_{split}'), eval_backbone=eval_backbone, post_resizer=post_resizer, max_items=n_samples, capture_all=True, capture_mean_cov=True)
    feats_real = stats_real.get_all_torch()
    feats_real = feats_real.detach().cpu().numpy().astype(np.float64)
    data_real_list.append(feats_real)
    data_real = copy.deepcopy(feats_real)

    # Load all synthetic distributions -> naive.
    gan_models = args.gan_models
    gan_steps = args.gan_steps
    gan_aval = os.listdir(samples_dir)
    gan_folders = [os.path.join(samples_dir, x, f'fake__{split}') for x in gan_aval if any(f'{gan_model}-train-' in x for gan_model in gan_models)]
    gan_folders = [os.path.join(x, f"step={y}") for x in gan_folders for y in gan_steps]
    pattern = fr"(?<={dataset_name}-)(\w+(?:-\w+)*)(?=-{split})"
    data_synthetic_list = []
    search_space = []
    with tqdm(total=len(gan_folders)) as pbar:
        for gan_sel in gan_folders:
            stats_synth = load_stats(cache_dir=cache_dir + gan_sel.split('samples')[-1], eval_backbone=eval_backbone, post_resizer=post_resizer,  max_items=n_samples, capture_all=True, capture_mean_cov=True)
            feats_synth = stats_synth.get_all_torch()
            feats_synth = feats_synth.detach().cpu().numpy().astype(np.float64)
            data_synthetic_list.append(feats_synth)
            search_space.append( re.search(pattern, gan_sel.split('/')[-3]).group(1) + '__' + gan_sel.split('/')[-1].split('=')[-1])
            pbar.update(1)
    data_synth = np.vstack(data_synthetic_list)

    # Load the dataframe containing the binary FID.
    try:
        df_intra = pd.read_excel(os.path.join(cache_dir, f'intra-{eval_backbone}-{post_resizer}-{n_samples}.xlsx'))
    except FileNotFoundError:
        pass

    # Dimensionality reduction using UMAP.
    if dim_reduction:

        # Check if data_real and data_synth already exist on the disk.
        flag_real = os.path.exists(os.path.join(os.path.join(source_dir,'ensemble'), f'data_real_umap-{eval_backbone}-{post_resizer}-{n_samples}.npy'))
        flag_synth = os.path.exists(os.path.join(os.path.join(source_dir,'ensemble'), f'data_synth_umap-{eval_backbone}-{post_resizer}-{n_samples}.npy'))
        if flag_real and flag_synth:
            print('Loading data_real and data_synth from the disk.')
            data_real = np.load(os.path.join(os.path.join(source_dir,'ensemble'), f'data_real_umap-{eval_backbone}-{post_resizer}-{n_samples}.npy'))
            data_synth = np.load(os.path.join(os.path.join(source_dir,'ensemble'), f'data_synth_umap-{eval_backbone}-{post_resizer}-{n_samples}.npy'))
        else:
            print('Reducing the dimensionality.')
            # First find the number of dimension to retain all the Variance.
            from sklearn.decomposition import PCA
            pca = PCA(n_components=0.95)
            pca.fit(data_real)
            n_components = pca.n_components_
            print(f"Number of components to retain 95% of the variance: {n_components}")

            # Plot the variance.
            import matplotlib.pyplot as plt
            plt.plot(np.cumsum(pca.explained_variance_ratio_))
            plt.xlabel('Number of components')
            plt.ylabel('Cumulative explained variance')
            # Plot a line definining the number of components to retain 95% of the variance.
            plt.axvline(x=n_components, color='r', linestyle='--')
            # Write in the legend the number of components to retain 95% of the variance.
            plt.legend(['Cumulative explained variance', f'Number of components to retain 95% of the variance: {n_components}'], loc='best')
            plt.tight_layout()
            plt.savefig(os.path.join(os.path.join(source_dir,'ensemble'), f'cumulative_explained_variance-{eval_backbone}-{post_resizer}-{n_samples}.png'), dpi=300, bbox_inches='tight')
            plt.show()

            # Apply UMAP.
            import umap
            reducer = umap.UMAP(n_components=n_components, random_state=42)
            data_real = reducer.fit_transform(data_real)
            data_synth = reducer.transform(data_synth)

            # Save to the disk.
            np.save(os.path.join(os.path.join(source_dir,'ensemble'), f'data_real_umap-{eval_backbone}-{post_resizer}-{n_samples}.npy'), data_real)
            np.save(os.path.join(os.path.join(source_dir,'ensemble'), f'data_synth_umap-{eval_backbone}-{post_resizer}-{n_samples}.npy'), data_synth)

        # Override data_real_list and data_synthetic_list.
        data_real_list = [data_real[i : i + n_samples] for i in range(0, n_samples * len(data_real_list), n_samples)]
        data_synthetic_list = [data_synth[i : i + n_samples] for i in  range(0, n_samples * len(data_synthetic_list), n_samples)]

    if synth_samples_reduction < n_samples:
        print('Reducing the number of synthetic samples.')
        # Reduce the number of samples for each distribution with random sampling.
        data_synthetic_list = [x[np.random.choice(x.shape[0], synth_samples_reduction, replace=False)] for x in data_synthetic_list]
        data_synth = np.vstack(data_synthetic_list)

    # Ensemble search.
    study = optuna.create_study(directions=directions, sampler=optuna.samplers.TPESampler(multivariate=True))

    func = lambda trial: util_ensemble.objective(
        trial=trial,
        samples_real=data_real,
        samples_naive_list=data_synthetic_list,
        search_space=search_space,
        name_obj=name_obj,
        pairwise=pairwise_intra,
        df_intra = df_intra
    )

    study.optimize(func, n_trials=n_trial)
    df = study.trials_dataframe(attrs=('number', 'value', 'params', 'state'))

    # Save the study ready to be uploaded.
    df = df[df['state'] == "COMPLETE"]
    df.reset_index(inplace=True, drop=True)
    df.to_excel(os.path.join(ensemble_dir, 'optuna_study.xlsx'), index=False)

    best_trials_idx, best_trials_name = util_ensemble.extract_best_trial(ensemble_dir=ensemble_dir, name_obj=name_obj, study=study)
    print('\n')
    print('Selected Ensemble.')
    print(best_trials_name)
    sel_idx = copy.deepcopy(best_trials_idx)
    sel_folders = [x for x in gan_folders if sel_idx[gan_folders.index(x)] == 1]

    # *
    # **
    # ***
    # ***
    # **
    # *
    # Downstream task.
    import yaml
    import torch
    import torch.nn as nn
    import torch.optim as optim
    from torch.optim import lr_scheduler
    import collections

    from src.data_util import Dataset_
    from src.general_utils.util_data import EnsembleDataset
    from src.cnn_models.models import ResNet18, ResNet50
    from src.general_utils import util_cnn

    with open(args.cfg_file) as file:
        cfg = yaml.load(file, Loader=yaml.FullLoader)

    init_w = args.init_w
    n_train = args.n_train
    img_size = cfg['DATA']['img_size']

    gpu_ids = int(args.gpu_ids)
    num_workers = args.num_workers

    # Device.
    if gpu_ids >= 0:
        # Check if available.
        if not torch.cuda.is_available():
            device = torch.device('cpu')
        else:
            device = torch.device('cuda:{}'.format(gpu_ids))
    else:
        device = torch.device('cpu')

    # Directories.
    filename = f'downstream_task_ensemble-n_trial__{n_trial}-name_obj__{name_obj}-pairwise_intra__{pairwise_intra}-dim_reduction__{dim_reduction}-synth_samples_reduction__{synth_samples_reduction}-{eval_backbone}-{post_resizer}-{n_samples}'
    #filename = f'{now.strftime("%Y-%m-%d_%H-%M-%S")}_{filename}'
    reports_dir = os.path.join(reports_dir, filename)
    util_path.create_dir(reports_dir)
    with open(os.path.join(reports_dir, f"ensemble.txt"), 'w') as f:
        f.write(f"\nEnsemble: {best_trials_name}")
        f.write("\n")
        f.write(f"\nFolder ensemble: {sel_folders}")

    datasets_real = {
        step: Dataset_(
                data_name=cfg['DATA']['name'],
                data_dir=cfg['DATA']['data_dir'],
                train=True if step == 'train' else False,
                split=step,
                crop_long_edge=cfg['PRE']['crop_long_edge'],
                resize_size=cfg['PRE']['resize_size'],
                resizer=cfg['PRE']['pre_resizer'],
                random_flip=cfg['PRE']['apply_rflip'],
                normalize=cfg['PRE']['normalize'],
                hdf5_path=os.path.join(cfg['DATA']['hdf5'], f'{dataset_name}_{img_size}_{step}.hdf5') if cfg['DATA']['hdf5'] is not None else None,
                load_data_in_memory=cfg['DATA']['load_data_in_memory'],
                cfgs=cfg) for step in ['train', 'val', 'test']
    }

    if dataset_name == 'AIforCOVID':
        classes = cfg['DATA']['classes']
        class_to_idx = {c: i for i, c in enumerate(sorted(classes))}
        idx_to_class = {i: c for c, i in class_to_idx.items()}

        weight = [
            len(datasets_real['train']) / (len(classes) * len(datasets_real['train'].labels[datasets_real['train'].labels == class_to_idx[c]])) for c in classes
        ]
    elif dataset_name in ['pneumoniamnist', 'breastmnist', 'retinamnist']:
        idx_to_class = datasets_real['test'].data.info['label']
        idx_to_class = {int(k): v for k, v in idx_to_class.items()}
        class_to_idx = {v: k for k, v in idx_to_class.items()}
        classes = [idx_to_class[i] for i in range(len(idx_to_class))]
        weight = [
            len(datasets_real['train'].data) / (len(classes) * len(datasets_real['train'].data.labels[datasets_real['train'].data.labels == class_to_idx[c]])) for c in classes
        ]
    else:
        raise NotImplementedError

    # Synthetic data.
    weights = util_ensemble.initialize_ensemble_weights(init_w=init_w, gan_list=sel_folders)
    dataset_synth = EnsembleDataset(folders=sel_folders, weights=weights)
    datasets = {
        'train': dataset_synth,
        'val': datasets_real['val'],
        'test': datasets_real['test']
    }

    data_loaders = {
        'train': torch.utils.data.DataLoader(datasets['train'], batch_size=cfg['TRAINER']['batch_size'], shuffle=True, num_workers=num_workers),
        'val': torch.utils.data.DataLoader(datasets['val'], batch_size=cfg['TRAINER']['batch_size'], shuffle=False, num_workers=num_workers),
        'test': torch.utils.data.DataLoader(datasets['test'], batch_size=cfg['TRAINER']['batch_size'], shuffle=False, num_workers=num_workers)
    }

    # List to save results for each training.
    results = collections.defaultdict(lambda: [])
    for idx_train in np.arange(n_train):
        results['Training'].append(idx_train)
        print(f"Training {idx_train + 1}/{n_train}...")

        print('==> Building and training model...')
        if cfg['MODEL']['name'] == 'resnet18':
            model = ResNet18(input_channels=cfg['DATA']['img_channels'], num_classes=cfg['DATA']['num_classes'])
        elif cfg['MODEL']['name'] == 'resnet50':
            model = ResNet50(input_channels=cfg['DATA']['img_channels'], num_classes=cfg['DATA']['num_classes'])
        else:
            raise NotImplementedError
        model = model.to(device)

        # Loss function.
        criterion = nn.CrossEntropyLoss(weight=torch.Tensor(weight)).to(device)
        optimizer = optim.Adam(model.parameters(), lr=cfg['TRAINER']['optimizer']['lr'])
        scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode=cfg['TRAINER']['scheduler']['mode'],  patience=cfg['TRAINER']['scheduler']['patience'])

        model, history = util_cnn.train_model(
            model=model,
            data_loaders=data_loaders,
            criterion=criterion,
            optimizer=optimizer,
            scheduler=scheduler,
            num_epochs=cfg['TRAINER']['max_epochs'],
            early_stopping=cfg['TRAINER']['early_stopping'],
            warmup_epoch=cfg['TRAINER']['warmup_epoch'],
            model_dir=reports_dir,
            device=device,
            n_samples=n_samples,
            to_disk=False
        )

        # Plot Training.
        util_cnn.plot_training(history=history, plot_training_dir=reports_dir, plot_name_loss=f'Loss_training_{idx_train}', plot_name_acc=f'Acc_training_{idx_train}')

        # Test model.
        test_results = util_cnn.evaluate(model=model, data_loader=data_loaders['test'], device=device, idx_to_class=idx_to_class)

        # Update report.
        results["ACC"].append(test_results['all'])
        for c in classes:
            results["ACC %s" % str(c)].append(test_results[c])
        results['recall'].append(test_results['recall'])
        results['precision'].append(test_results['precision'])
        results['specificity'].append(test_results['specificity'])
        results['f1_score'].append(test_results['f1_score'])
        results['g_mean'].append(test_results['g_mean'])

        # Save results
        reports_file = os.path.join(reports_dir, f'results_training_{idx_train}.xlsx')
        results_frame = pd.DataFrame(dict([(k, pd.Series(v)) for k, v in results.items()]))
        results_frame.to_excel(reports_file, index=False)

    # Compute mean and std.
    results_frame = pd.DataFrame(dict([(k, pd.Series(v)) for k, v in results.items()]))
    results_frame.loc['mean'] = results_frame.mean()
    results_frame.loc['std'] = results_frame.std()

    # Save Results
    reports_file = os.path.join(reports_dir, 'results.xlsx')
    results_frame.to_excel(reports_file, index=True)

    print("May the force be with you.")


"""
This module defines all methods which deal with the contributing files
"""
import os

from github import Github
from github.GithubException import UnknownObjectException

g = Github(os.environ['GITHUB_GPG_KEY'])

def get_contributing_files(repo_addresses: list[str]) -> dict[str, str]:
    """Method gets multiple contributing files from git and decodes them in utf-8 form.
    Then it adds these files to a dictionairy

    :param repo_addresses: repository addresses in a list formatted as ['author1/name1', 'author2/name2'...]
    :type repo_addresses: list[str]
    :return: A dictionairy from a repository address to its contributing file
    :rtype dict[str, str]
    """
    contributing_files = {}
    for address in repo_addresses:
        repo = g.get_repo(address)
        try:
            contributing_contents = repo.get_contents("CONTRIBUTING.md").decoded_content.decode(
                'utf-8')  # decode to utf-8
        except UnknownObjectException:
            contributing_contents = None
        contributing_files[address] = contributing_contents
    return contributing_files


import requests
import os
from requests.auth import HTTPBasicAuth
import yaml
from dotenv import load_dotenv

load_dotenv()

access_token = os.getenv('access_token')

def read_yaml_from_github(access_token):
    # GitHub API endpoint for raw content
    raw_url = f'https://raw.githubusercontent.com/hmrc/cip-data-domain-mapping/main/data-domain-to-audit-source-mapping.yaml'

    # Make a GET request with the personal access token for authentication
    response = requests.get(raw_url, auth=HTTPBasicAuth('token', access_token))

    # Check if the request was successful (status code 200)
    if response.status_code == 200:
        # Parse the YAML content
        yaml_content = yaml.safe_load(response.text)
        return yaml_content
    else:
        print(f"Failed to fetch YAML file. Status code: {response.status_code}")
        return None

# Replace these values with your repository details and personal access token

# Read YAML content from the private GitHub repository
yaml_data = read_yaml_from_github(access_token)
audit_source = 'income-tax-subscription-eligibility;submit-vat-return-frontend;import-control-entry-declaration-outcome;fake-audit-source'

def find_data_domains(yaml_data, audit_source):

    if yaml_data:
        print({"Response":"Successfully obtained yaml data from repo"})

    audit_source = audit_source.lower()
    audit_source = audit_source.replace('_', '-')

    audit_source_list = audit_source.split(';')
    data_domain_list = []
    found_audit_sources = []

    for key, values in yaml_data.items():
        for audit_sources in audit_source_list:
            if audit_sources in values:

                data_domain_list.append(key)
                found_audit_sources.append(audit_sources)
                print(f'{key}: {audit_sources}')

            # else:
            #     data_domain_list.append(f'{audit_sources} not found')

    missing = list(sorted(set(audit_source_list) - set(found_audit_sources)))
    data_domains = list(set(data_domain_list))

    return data_domains, missing

if __name__=='__main__':
    yaml_data2 = read_yaml_from_github(access_token)
    data_doms = find_data_domains(yaml_data2, audit_source)
    print(data_doms)


import torch
import numpy as np
import matplotlib.pyplot as plt
import os
import sklearn as sk
from Code2 import FeatureEngineering as FE
import torchvision
import tqdm
import collections

def to_cuda(elements):
    if torch.cuda.is_available():
        if type(elements) == tuple or type(elements) == list:
            return [x.cuda() for x in elements]
        return elements.cuda()
    return elements

class FullyConnectedModel(torch.nn.Module):

    def __init__(self):
        super().__init__()
        # Number of input nodes (size of images)
        num_input_nodes = 20*20

        hidden_input = 32
        # Number of classes in the chars74k-lite dataset
        num_classes = 26

        # Define our model
        self.classifier = torch.nn.Sequential(
            torch.nn.Linear(num_input_nodes, hidden_input),
            torch.nn.ReLU(),
            torch.nn.Linear(hidden_input, num_classes),
            torch.nn.ReLU()
        )

    def forward(self, x):
        # Runs a forward pass on the images
        x = x.view(-1, 20*20)
        out = self.classifier(x.float())
        return out

class Trainer:

    def __init__(self,
                 model,
                 dataloader_train,
                 dataloader_val,
                 batch_size,
                 loss_function,
                 optimizer):
        self.dataloader_train = dataloader_train
        self.dataloader_val = dataloader_val
        self.batch_size = batch_size

        self.model = model
        self.loss_function = loss_function
        self.optimizer = optimizer

    def train(self, num_epochs):
        tracked_train_loss = collections.OrderedDict()
        tracked_test_loss = collections.OrderedDict()
        global_step = 0

        for epoch in range(num_epochs):
            avg_loss = 0

            for batch_it, (images, target) in enumerate(
                    tqdm.tqdm(self.dataloader_train,
                              desc=f"Training epoch {epoch}")):
                # images has shape: [batch_size, 1, 20, 20]
                # target has shape: [batch_size]

                batch_it, (images, target) = to_cuda([batch_it, (images, target)])

                # Perform forward pass
                logits = self.model(images)

                # Compute loss
                loss = self.loss_function(logits, target)

                avg_loss += loss.detach().item()
                # Perform backpropagation
                loss.backward()

                # Update our parameters with gradient descent
                self.optimizer.step()

                # Reset our model parameter gradients to 0
                self.optimizer.zero_grad()

                # Track the average loss for every 500th image
                if batch_it % (500//self.batch_size) == 0 and batch_it != 0:
                    avg_loss /= (500//self.batch_size)
                    tracked_train_loss[global_step] = avg_loss
                    avg_loss = 0
                global_step += self.batch_size
            # Compute loss and accuracy on the test set
            test_loss, test_acc = compute_loss_and_accuracy(
                self.dataloader_val, self.model, self.loss_function
            )
            tracked_test_loss[global_step] = test_loss
        return tracked_train_loss, tracked_test_loss

def plot_loss(loss_dict, label):
    global_steps = list(loss_dict.keys())
    loss = list(loss_dict.values())
    plt.plot(global_steps, loss, label=label)

def compute_loss_and_accuracy(dataloader, model, loss_function):
    """
    Computes the total loss and accuracy over the whole dataloader
    Args:
        dataloder: Test dataloader
        model: torch.nn.Module
        loss_function: The loss criterion, e.g: nn.CrossEntropyLoss()
    Returns:
        [loss_avg, accuracy]: both scalar.
    """
    model.eval()
    # Tracking variables
    loss_avg = 0
    total_correct = 0
    total_images = 0
    total_steps = 0
    with torch.no_grad():  # No need to compute gradient when testing
        for (X_batch, Y_batch) in dataloader:
            # Forward pass the images through our model
            output_probs = model(X_batch)
            # Compute loss
            loss = loss_function(output_probs, Y_batch)

            # Predicted class is the max index over the column dimension
            predictions = output_probs.argmax(dim=1).squeeze()
            Y_batch = Y_batch.squeeze()

            # Update tracking variables
            loss_avg += loss.item()
            total_steps += 1
            total_correct += (predictions == Y_batch).sum().item()
            #print("Prediction:")
            #print(predictions,Y_batch)
            total_images += predictions.shape[0]
    model.train()
    loss_avg = loss_avg / total_steps
    accuracy = total_correct / total_images
    return loss_avg, accuracy

class MyData(torch.utils.data.Dataset):

    def __init__(self,set):
        self.samples = []
        for i,elements in enumerate(set):
            for element in elements:
                self.samples += [(element,i)]

    def __getitem__(self,idx):
        return self.samples[idx]

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

def load_dataset(batch_size, fe, image_transform):

    dataset_train = MyData(fe.training_set)

    dataset_test = MyData(fe.test_set)

    dataloader_train = torch.utils.data.DataLoader(
        dataset_train,
        batch_size=batch_size,
        shuffle=True,
    )
    dataloader_test = torch.utils.data.DataLoader(
        dataset_test,
        batch_size=batch_size,
        shuffle=False
    )
    return dataloader_train, dataloader_test

if __name__ =="__main__":

    # ### Hyperparameters & Loss function

    # Hyperparameters
    batch_size = 32
    learning_rate = .11
    num_epochs = 5

    fe = FE()
    fe.HOG()
    fe.scaling()
    fe.splitDataset()

    # Use CrossEntropyLoss for multi-class classification
    loss_function = torch.nn.CrossEntropyLoss()

    # Model definition
    model = to_cuda(FullyConnectedModel())

    # Define optimizer (Stochastic Gradient Descent)
    optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

    image_transform = torchvision.transforms.Compose([
      torchvision.transforms.ToTensor()

    ])
    dataloader_train, dataloader_val = load_dataset(batch_size, fe, image_transform=image_transform)

    trainer = Trainer(
      model=model,
      dataloader_train=dataloader_train,
      dataloader_val=dataloader_val,
      batch_size=batch_size,
      loss_function=loss_function,
      optimizer=optimizer
    )
    train_loss_dict, val_loss_dict = trainer.train(num_epochs)

    final_loss, final_acc = compute_loss_and_accuracy(dataloader_val, model, loss_function)
    print(f"Final Test Cross Entropy Loss: {final_loss}. Final Test accuracy: {final_acc}")


import sys
from enum import Enum
import cv2
import matplotlib.pyplot as plt
from .setting import *

class DectionMethod(Enum):
    ColorFilter = 1

class ProcessType(Enum):
    CenterCoordinate = 1
    Bound = 2

###########################################
class CamerSystem:
    # 网络摄像头URL示例 r"http://192.168.0.120:8080/?action=stream"
    cameraMainURL = r"C:\Users\Arnoliu\PycharmProjects\opecvtry\toolbox\demoVideo.mp4"
    cameraViceURL = 0

    def __init__(self, oneOrTwoCamera=1):
        # 根据用户的选择开启几个摄像头
        self.CameraNumber = oneOrTwoCamera
        assert oneOrTwoCamera != 1 or oneOrTwoCamera != 2
        if oneOrTwoCamera >= 1:
            self.cap_l = cv2.VideoCapture(CamerSystem.cameraMainURL)
        if oneOrTwoCamera == 2:
            self.cap_r = cv2.VideoCapture(CamerSystem.cameraViceURL)

    # 读取配置文件
    def loadConfigue(self):
        self.lower_red, self.upper_red = readColorIni()

    def takeImg(self):
        ret, frame = self.cap_l.read()
        if not ret:
            print("No Stream or stream stopped!")
            return None
        return frame

    def getFPS(self):
        return self.cap_l.get(propId=cv2.CAP_PROP_FPS)

    def releasCam(self):
        if self.CameraNumber >= 1:
            self.cap_l.release()
        if self.CameraNumber == 2:
            self.cap_r.release()

    # MotionDect
    # 有固定的背景https://www.yyearth.com/article/18-09/242.html

    fgbg = cv2.createBackgroundSubtractorKNN(history=100, detectShadows=False)

    def MotionThreshold(self, img=None, keepOrigin=False):
        if img is None:
            frame = self.takeImg()
            if frame is None:
                return None
        else:
            frame = img
        blured = cv2.GaussianBlur(frame, (3, 3), 0)
        knersize, iterations = readMotionIni()

        fgmask = self.fgbg.apply(blured)
        fgmask = open_mor(fgmask, knersize, iterations)
        if keepOrigin:
            return frame, fgmask
        return fgmask

    def ColorThreshold(self, img=None, blurKernelSize=5, colorBound=None, showSource=True, keepOrigin=False):
        self.loadConfigue()
        if colorBound == None:
            lower_bound = self.lower_red
            upper_bound = self.upper_red
        else:
            lower_bound = colorBound[0]
            upper_bound = colorBound[1]

        if img is None:
            frame = self.takeImg()
            if frame is None:
                return None
        else:
            frame = img
        if showSource:
            cv2.imshow("source", frame)
        blured = cv2.GaussianBlur(frame, (blurKernelSize, blurKernelSize), 0)
        color_fliter = cv2.inRange(blured, lower_bound, upper_bound)
        color_fliter = cv2.erode(color_fliter, (10, 10), iterations=3)
        color_fliter = open_mor(color_fliter, 5, 1)

        if keepOrigin:
            return frame, color_fliter
        return color_fliter

    matchImg = cv2.imread("obj.jpg")

    def GetCenterByMatchTemplet(self, img=None,method=1):
        if img is None:
            frame = self.takeImg()
            if frame is None:
                return -1,-1
        else:
            frame = img
        match = cv2.matchTemplate(frame, self.matchImg, method=method)
        min_val, max_val, min_loc, max_loc =cv2.minMaxLoc(match)
        # If the method is TM_SQDIFF or TM_SQDIFF_NORMED, take minimum
        if method in [0, 1]:
            top_left = min_loc
        else:
            top_left = max_loc
        drawPointonImg(frame,top_left[0],top_left[1])
        return top_left

    def GetCenterByLongestContour(self, img, type=ProcessType.CenterCoordinate, debug=False):
        if sys.platform.startswith('win'):
            contours_m, hierarchy_m = cv2.findContours(img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        else:
            h, contours_m, hierarchy_m = cv2.findContours(img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        areas = [cv2.contourArea(c) for c in contours_m]
        if len(areas) == 0:
            return (-1, -1, -1, -1)
        longestArea = contours_m[np.nanargmax(areas)]

        if type == ProcessType.CenterCoordinate:
            hull = cv2.convexHull(longestArea)

            if debug:
                debugImg = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
                length = len(hull)
                for i in range(len(hull)):
                    cv2.line(debugImg, tuple(hull[i][0]), tuple(hull[(i + 1) % length][0]), (0, 255, 0), 2)
                cv2.imshow("Debug", debugImg)
            return hull.mean(axis=0).tolist()[0]
            # return round((2 * x + w) / 2), round((2 * y + h) / 2)
        else:
            (x, y, w, h) = cv2.boundingRect(longestArea)
            return x, y, w, h

# region Open And Close
def open_mor(src, kernelsize, iter):
    kernel = np.ones((kernelsize, kernelsize), np.uint8)
    opening = cv2.morphologyEx(src, cv2.MORPH_OPEN, kernel, iterations=iter)
    return opening

# ! Do not use this
def close_mor(src, kernelsize: int, iter):
    kernel = np.ones((kernelsize, kernelsize), np.uint8)
    opening = cv2.morphologyEx(src, cv2.MORPH_CLOSE, kernel, iterations=iter)
    return opening

# endregion
# 直接计算图像的重心
def centerPoint(thimg):
    mom = cv2.moments(thimg)
    if mom["m00"] == 0:
        return 0, 0
    center_x = round(mom["m10"] / mom["m00"])
    center_y = round(mom["m01"] / mom["m00"])
    return center_x, center_y

# Draw Tool
def drawPointonImg(img, px, py):
    if img.shape[2]!=3:
        img_rgb = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
    else:
        img_rgb=img.copy()
    cv2.circle(img_rgb, (px, py), 3, (0, 0, 255))
    cv2.imshow("centerPoint", img_rgb)

def drawCenterPoint(img, imgName):
    meanX, meanY = centerPoint(img)
    img_rgb = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
    cv2.circle(img_rgb, (meanX, meanY), 3, (0, 0, 255))
    cv2.imshow("centerPoint" + str(imgName), img_rgb)
    return meanX, meanY

def drawCont(thimg):
    tempimg = thimg.copy()
    contours = cv2.findContours(tempimg)
    img_rgb = cv2.cvtColor(tempimg, cv2.COLOR_GRAY2BGR)
    for point in contours:
        cv2.circle(tempimg, point, 1, (0, 0, 255))
    cv2.imshow("Cont", tempimg)


import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.patches import Rectangle

from .plotting_disease import setup_x_axis

mpl.rcParams['xtick.major.pad'] = '8'
mpl.rcParams['ytick.major.pad'] = '8'

try:
    import brewer2mpl as brew

    colors = ['#8dd3c7', '#ffffb3', '#bebada', '#fb8072', '#80b1d3', '#fdb462', '#b3de69']
    #colors = brew.get_map('Pastel2', 'qualitative', 7).mpl_colors
    #colors = ['k'] + brew.get_map('Pastel2', 'qualitative', 7).mpl_colors
    #colors[6] = 'c' # replace invisible yellow
    rb_colors = ['k'] + brew.get_map('YlOrRd', 'sequential', 9).mpl_colors
except ImportError:
    colors = ['k', 'r', 'b', 'g', 'm', 'c', '0.5']

## -------------------------------------------------------------------------------------------
## from http://www.sron.nl/~pault/
## -------------------------------------------------------------------------------------------

# colour table in HTML hex format
hexcols = ['#332288', '#88CCEE', '#44AA99', '#117733', '#999933', '#DDCC77',
           '#CC6677', '#882255', '#AA4499', '#661100', '#6699CC', '#AA4466',
           '#4477AA']

greysafecols = ['#809BC8', '#FF6666', '#FFCC66', '#64C204']

xarr = [[12],
        [12, 6],
        [12, 6, 5],
        [12, 6, 5, 3],
        [0, 1, 3, 5, 6],
        [0, 1, 3, 5, 6, 8],
        [0, 1, 2, 3, 5, 6, 8],
        [0, 1, 2, 3, 4, 5, 6, 8],
        [0, 1, 2, 3, 4, 5, 6, 7, 8],
        [0, 1, 2, 3, 4, 5, 9, 6, 7, 8],
        [0, 10, 1, 2, 3, 4, 5, 9, 6, 7, 8],
        [0, 10, 1, 2, 3, 4, 5, 9, 6, 11, 7, 8]]

# get specified nr of distinct colours in HTML hex format.
# in: nr - number of colours [1..12]
# returns: list of distinct colours in HTML hex
def get_distinct(nr, bw=False, br=False):
    # check if nr is in correct range

    if nr <= 2 and bw:
        return ['#333333', '#AAAAAA']

    if nr <= 2 and br:
        print("BBBBBBB")
        return ['#FFFFFF', '#FF0000']

#    if nr <= 2 and br:
#        return ['#333333', '#CC6677']

    if nr < 1 or nr > 12:
        print("wrong nr of distinct colours!")
        return
    # get list of indices
    lst = xarr[nr - 1]
    # generate colour list by stepping through indices and looking them up
    # in the colour table
    i_col = 0
    col = [0] * nr
    for idx in lst:
        col[i_col] = hexcols[idx]
        i_col += 1
    return col

## -------------------------------------------------------------------------------------------

# ideally this should be expanded, and be settable by main.py
class lp(object):
    lw = 1
    fs = 16 # 16 is very big (John thinks 12 looks better)

def plot_lines(ax, x_values, y_values, x_label, y_label, y_errs=None, series_labels=None, legend_title=None,
               legend_loc='upper left', linestyle='-', lineweight=None, bw=False, br=False, reversed=False, xmin=None):
    """
    Plotting function for lines, with or without error bars.
    """

    # if data values submitted as a single series, wrap in list
    if not (isinstance(y_values[0], list)
            or isinstance(y_values[0], np.ndarray)
            or isinstance(y_values[0], tuple)):
        y_values = [y_values]

    #print len(x_values), len(y_values[0])
    assert len(x_values) == len(y_values[0])

    if len(y_values) <= 12:
        loc_colors = get_distinct(len(y_values), bw, br)
    else:
        # get max number of colours
        loc_colors = get_distinct(12)

    # same with error values
    if y_errs is not None and (not isinstance(y_errs[0], (list, np.ndarray, tuple))):
        y_errs = [y_errs]

    # expand y_errs and series_labels if not present
    if not y_errs:
        y_errs = [None] * len(y_values)
    if not series_labels:
        series_labels = [None] * len(y_values)

    if reversed:
        y_values.reverse()
        loc_colors.reverse()

    # plot lines
    for i, (cur_values, cur_errs, cur_label) in enumerate(zip(y_values, y_errs, series_labels)):
        #print bw, cur_label, loc_colors[i]
        if cur_errs is not None:
            ax.errorbar(x_values, cur_values, yerr=cur_errs, label=cur_label, color=loc_colors[i],
                        lw=lp.lw if not lineweight else lineweight, ls=linestyle)
        else:
            #print len(x_values), len(cur_values)
            ax.plot(x_values, cur_values, label=cur_label, color=loc_colors[i % len(loc_colors)],
                    lw=lp.lw if not lineweight else lineweight, ls=linestyle)

    # set labels and legend
    ax.set_xlabel(x_label, fontsize=lp.fs)
    ax.set_ylabel(y_label, fontsize=lp.fs)

    if xmin is not None:
        ax.set_xlim(xmin=xmin)

    if series_labels[0]:
        leg = ax.legend(scatterpoints=1, numpoints=1, ncol=2,
                        title=legend_title, labelspacing=0.1, columnspacing=1,
                        loc=legend_loc)
        plt.setp(leg.get_texts(), fontsize='small')
        leg.get_frame().set_alpha(0.0)

    plt.setp(ax.get_xticklabels(), fontsize=lp.fs)
    plt.setp(ax.get_yticklabels(), fontsize=lp.fs)

def plot_points_lin_reg(ax, times, y_values, x_label, y_label, series_labels=None, y_min=None, y_max=None,
                        x_offset=None, align=None,
                        legend_title=None, legend_loc='upper left', plot_title=None, bw=False):
    """
    Plot a set of data points together with a line of best fit.
    """
    if not (isinstance(y_values[0], (list, np.ndarray, tuple))):
        y_values = [y_values]

    if not series_labels:
        series_labels = [None] * len(y_values)

    loc_colors = get_distinct(len(y_values), bw)

    if align == 'centre':
        offset = (times['et'] - times['st']) / (len(y_values[0]) + 1) / 2.0
        #print times['st'], times['et'], offset
        x_values = np.linspace(times['st'] + offset, times['et'] - offset, len(y_values[0])) / times['t_per_year']
        #print xvals
    else:
        x_values = np.linspace(times['st'], times['et'], len(y_values[0])) / times['t_per_year']
        #x_values = np.linspace(times['st'], times['et'], len(y_values[0])) / times['t_per_year']

    for i, (cur_values, cur_label) in enumerate(zip(y_values, series_labels)):
        p = np.polyfit(x_values[~np.isnan(cur_values)], cur_values[~np.isnan(cur_values)], 1)
        #print cur_values
        ax.plot(x_values, cur_values, color=loc_colors[i], marker='x', lw=0)
        ax.plot(x_values, p[0] * x_values + p[1], label=cur_label, color=loc_colors[i], linestyle='-', lw=lp.lw)

    if plot_title:
        ax.set_title(plot_title, fontsize=lp.fs)
    ax.set_xlabel(x_label, fontsize=lp.fs)
    ax.set_ylabel(y_label, fontsize=lp.fs)
    ncol = 2 if len(y_values) > 5 else 1
    if series_labels[0]:
        leg = ax.legend(scatterpoints=1, numpoints=1, ncol=ncol, title=legend_title, labelspacing=0.1, columnspacing=1,
                        loc=legend_loc)
        plt.setp(leg.get_texts(), fontsize='small')
        leg.get_frame().set_alpha(0.0)

    if y_min is not None:
        ax.set_ylim(ymin=y_min)
    if y_max is not None:
        ax.set_ylim(ymax=y_max)

    plt.setp(ax.get_xticklabels(), fontsize=lp.fs)
    plt.setp(ax.get_yticklabels(), fontsize=lp.fs)

    setup_x_axis(ax, times, x_offset, set_xlim=False)

def plot_stacked_area(ax, x_values, y_values, x_label, y_label, series_labels=None,
                      legend_title=None, plot_title=None, bw=False):
    if not series_labels:
        series_labels = [None] * len(y_values)

    #print series_labels

    p = []

    colors = get_distinct(len(y_values), bw)

    ax.fill_between(x_values, 0, y_values[0], lw=0, edgecolor='grey', facecolor=colors[0])
    p.append(Rectangle((0, 0), 1, 1, lw=0, edgecolor='grey', facecolor=colors[0]))  # for legend
    for i, (cur_bottom, cur_top) in enumerate(zip(y_values[:-1], y_values[1:])):
        ax.fill_between(x_values, cur_bottom, cur_top, lw=0, edgecolor='grey', facecolor=colors[i + 1])
        p.append(Rectangle((0, 0), 1, 1, lw=0, edgecolor='grey', facecolor=colors[i + 1]))

    ax.set_ylim(ymax=1.0)

    # set labels and legend
    if plot_title:
        ax.set_title(plot_title, fontsize=lp.fs)
    ax.set_xlabel(x_label, fontsize=lp.fs)
    ax.set_ylabel(y_label, fontsize=lp.fs)
    if series_labels[0]:
        leg = ax.legend(p, series_labels, ncol=2, title=legend_title, labelspacing=0.1, columnspacing=1,
                        loc='upper right')
        plt.setp(leg.get_texts(), fontsize='small')
        leg.get_frame().set_alpha(0.8)
        leg.get_frame().set_linewidth(0.0)
        #leg = ax.legend(p, series_labels)

    plt.setp(ax.get_xticklabels(), fontsize=lp.fs)
    plt.setp(ax.get_yticklabels(), fontsize=lp.fs)

def output_bars(ofile, x_values, y_values, x_label, y_label, x_ticklabels=None, y_errs=None, ymin=0.0, ymax=None,
                series_labels=None,
                legend_loc=2, rotate_ticks=False, plot_title=None, bw=False):
    fig = plt.figure()
    ax = fig.add_subplot(111)

    plot_bars(ax, x_values, y_values, x_label, y_label, x_ticklabels, y_errs, ymin, ymax, series_labels,
              legend_loc, rotate_ticks, plot_title, bw)

    fig.tight_layout()
    fig.savefig(ofile)

def plot_bars(ax, x_values, y_values, x_label, y_label, x_ticklabels=None, y_errs=None, ymin=0.0, ymax=None,
              series_labels=None,
              legend_loc=2, rotate_ticks=False, plot_title=None, bw=False):
    # if data values submitted as a single series, wrap in list
    if not (isinstance(y_values[0], list)
            or isinstance(y_values[0], np.ndarray)
            or isinstance(y_values[0], tuple)):
        y_values = [y_values]

    #print len(x_values), len(y_values[0])
    assert len(x_values) == len(y_values[0])

    loc_colors = get_distinct(len(y_values), bw)

    # same with error values
    if y_errs is not None and (not isinstance(y_errs[0], (list, np.ndarray, tuple))):
        y_errs = [y_errs]

    # expand y_errs and series_labels if not present
    if not y_errs:
        y_errs = [None] * len(y_values)
    if not series_labels:
        series_labels = [None] * len(y_values)

    width = 0.8 / len(y_values)

    x_values = np.array(x_values)

    # plot bars
    for i, (cur_values, cur_errs, cur_label) in enumerate(zip(y_values, y_errs, series_labels)):
        ax.bar(x_values + i * width, cur_values, width, lw=0, edgecolor='grey', color=loc_colors[i], ecolor='dimgrey',
               align='center', yerr=cur_errs, label=cur_label)

    if plot_title:
        ax.set_title(plot_title, fontsize=lp.fs)
    ax.set_ylabel(y_label, fontsize=lp.fs)
    ax.set_xlabel(x_label, fontsize=lp.fs)
    if len(y_values) == 1:
        ax.set_xlim(xmin=-1, xmax=len(x_values))
    else:
        ax.set_xlim(xmin=-0.5, xmax=len(x_values))

    ax.set_ylim(ymin=ymin)
    if ymax is not None:
        ax.set_ylim(ymax=ymax)

    if x_ticklabels:
        ax.set_xticks(np.array(x_values) + len(y_values) / 2 * width)
        ax.set_xticklabels(tuple(x_ticklabels), rotation=(20 if rotate_ticks else 0))

    ax.xaxis.set_ticks_position('none')

    if series_labels[0]:
        leg = ax.legend(loc=legend_loc)
        leg.get_frame().set_alpha(0.0)

    plt.setp(ax.get_xticklabels(), fontsize=lp.fs)
    plt.setp(ax.get_yticklabels(), fontsize=lp.fs)

def plot_statistics(ofile, x_values, y_values, x_label, y_label, y_errs=None, series_labels=None,
                    ymin=None, xmax=None, xlogscale=False, ylogscale=False, title=None, bw=False,
                    legend_loc=None, legend_title=None):
    # create figure
    fig = plt.figure()
    ax = fig.add_subplot(111)

    plot_lines(ax, x_values, y_values, x_label, y_label, y_errs, series_labels, bw=bw,
               legend_loc=legend_loc, legend_title=legend_title)

    if xlogscale:
        ax.set_xscale('log')
    if ylogscale:
        ax.set_yscale('log')

    if ymin:
        if isinstance(ymin, (int, float)):
            ax.set_ylim(ymin=ymin)
        elif ymin == 'zero':
            ax.set_ylim(ymin=0)

    if title:
        ax.set_title(title, fontsize=lp.fs)

    fig.tight_layout()
    fig.savefig(ofile)


from imutils.perspective import four_point_transform
import imutils
import cv2
import numpy as np

def Get_Outline(image):
    #image = cv2.imread(input_dir)
    #image = np.uint8(image)
    ###image = np.array(image, dtype=np.uint8)
    ###image = (image*255).astype(np.uint8)
    #gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    blurred = cv2.GaussianBlur(image, (5,5),0)
    edged = cv2.Canny(blurred,75,200)
    return image,edged

def Get_cnt(edged):
    cnts = cv2.findContours(edged.copy(), cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_SIMPLE)
    cnts = cnts[1] if  imutils.is_cv3()  else   cnts[0]
    docCnt =None
    l1=[]
    nombre=0
    if len(cnts) > 0:
        cnts =sorted(cnts,key=cv2.contourArea,reverse=True)
        for c in cnts:
            peri = cv2.arcLength(c,True)                   # 輪廓按大小降序排序
            approx = cv2.approxPolyDP(c,0.02 * peri,True)  # 獲取近似的輪廓
            #if len(approx) ==4:                            # 近似輪廓有四個頂點
            #l1=[]
            while len(approx)==4:
                nombre+=1
                #print(len(approx))
                docCnt = approx
                '''
                docCnt = np.reshape(docCnt,(4,1,2))
                docCnt = np.reshape(docCnt,(-1,2))
                '''
                #print(docCnt)
                docCnt=docCnt.tolist()
                #print("-----------------")
                #docCnt = np.reshape(docCnt, (2, 4))
                #docCnt = docCnt.flatten()
                #print(docCnt)
                '''
                if nombre==1:
                    l1=docCnt
                else:
                    #print(docCnt)
                    print("-----------------")
                    #l1=np.array([l1,docCnt])
                '''
                l1.append(docCnt)

                break
    l1=np.array(l1)
    #print(l1)
    #print(nombre)
    #l1="sample"

    l1 = np.reshape(l1,(nombre,4,2))
    l1 = np.reshape(l1,(-1,2))
    #print(l1)
    '''
    print(l1[0][0])
    print(l1[0][1])
    print(max(l1[::3]))
    '''
    return l1,nombre

def region_cal(input, input_type = 'image'):
    if input_type == 'image':
        image = cv2.imread(input)
    else:
        image = input
    image,edged = Get_Outline(image)
    l1 = Get_cnt(edged)
    #print(l1)
    #print(docCnt)
    return l1

if __name__=="__main__":

    input = "./14_0.png"
    #region_cal(input_dir)
    region_cal(input)
    #region_cal("./detelet_vh.png")
#convex_hull


#!/usr/bin/python

import numpy as np
import matplotlib.pyplot as plt
import os

def amplitude(a):
	a							=	(a.imag**2+a.real**2)**0.5
	return a

def mean(a):
	mean_a						=	np.zeros(256)
	for i in range(256):
		for j in range(len(a)):
			mean_a[i]			+=	amplitude(a[j][i])
		mean_a[i]				=	mean_a[i]/len(a)
	return mean_a

def reject_matrix(XX_m):
	R		=	np.zeros(256)
	for i in range(256):
		if(np.abs(amplitude(XX_m[i])-np.mean(XX_m))>np.std(XX_m)*n_s):
			R[i]	=	1
	return R

def reject(R,corr):
	for j in range(256):
		if(R[j]==1):
			for i in range(i_n):
				corr[i][j]=0

def corr_read(s1,fopen,XX_1):
	j	=	0
	i	=	0
	while(i<int(s1)):
		p		=	fopen.read(1)
		ctr1	=	-1
		ctr2	=	-1
		e_flag	=	0
		while(p!='(' or p=='0' and i<int(s1)):
			if(p=='0'):
				#print p
				XX_1.append(complex(0,0))
				i	+=	2
				fopen.read(2)
			if(not(i<int(s1))):
				break
			p	=	fopen.read(1)
			i	+=	1
		if(not(i<int(s1))):
			break
		while((p!='+' or e_flag) and (p!='-' or ctr1==0 or e_flag)):
			p	=	fopen.read(1)
			ctr1	+=	1
			if(not(e_flag-ctr1)):
				e_flag	=	0
			if(p=='e'):
				e_flag	=	ctr1+2
		while(p!='j'):
			p	=	fopen.read(1)
			ctr2	+=	1
		fopen.seek(-ctr1-ctr2-2,1)
		r	=	float(fopen.read(ctr1))
		fopen.read(1)
		im	=	float(fopen.read(ctr2))
		#print complex(r,im)
		XX_1.append(complex(r,im))
		fopen.read(3)
		i	+=	ctr1+ctr2+5
		j	+=	1

file_name	=	'ch00_CASA_20170602_234118_000_cross_corr.txt'
#file_name	=	raw_input('File: ')
n_s		=	2
fopen	=	open(file_name,'r')
size	=	os.path.getsize(file_name)
s1		=	float(size)/4
ARR		=	[]

corr_read(size,fopen,ARR)

l	= int(float(len(ARR)))

n	=	int(l/4)
i_n	=	int(n/256)

ARR_256	=	np.zeros((4,i_n,256),dtype='complex')
XX		=	np.zeros((i_n,256),dtype='complex')
XY		=	np.zeros((i_n,256),dtype='complex')
YX		=	np.zeros((i_n,256),dtype='complex')
YY		=	np.zeros((i_n,256),dtype='complex')

for k in range(4):
	for i in range(i_n):
		for j in range(256):
			ARR_256[k][i][j]	=	ARR[(k*n)+(i*256)+j]

XX	=	ARR_256[0]
XY	=	ARR_256[1]
YY	=	ARR_256[2]
YX	=	ARR_256[3]

XX_m	=	mean(XX)
XY_m	=	mean(XY)
YY_m	=	mean(YY)
YX_m	=	mean(YX)

R_XX	=	reject_matrix(XX_m)
R_XY	=	reject_matrix(XY_m)
R_YY	=	reject_matrix(YY_m)
R_YX	=	reject_matrix(YX_m)

reject(R_XX,XX)
reject(R_XY,XY)
reject(R_YY,YY)
reject(R_YX,YX)

fig1	=	plt.figure(1)
plt.subplot(221)
plt.imshow(amplitude(XX.T),interpolation='bilinear',aspect='auto')
plt.ylabel('XX')
plt.subplot(222)
plt.imshow(amplitude(XY.T),interpolation='bilinear',aspect='auto')
plt.ylabel('XY')
plt.subplot(223)
plt.imshow(amplitude(YY.T),interpolation='bilinear',aspect='auto')
plt.ylabel('YY')
plt.subplot(224)
plt.imshow(amplitude(YX.T),interpolation='bilinear',aspect='auto')
plt.ylabel('YX')
fig1.savefig(file_name[:-4]+'_im(clean).png')

fig11	=	plt.figure(11)
plt.subplot(221)
plt.plot(amplitude(XX_m))
plt.plot([np.mean(XX_m)-n_s*np.std(XX_m) for i in range(len(XX_m))])
plt.plot([np.mean(XX_m)+n_s*np.std(XX_m) for i in range(len(XX_m))])
plt.ylabel('XX')
plt.subplot(222)
plt.plot(amplitude(XY_m))
plt.plot([np.mean(XY_m)-n_s*np.std(XY_m) for i in range(len(XY_m))])
plt.plot([np.mean(XY_m)+n_s*np.std(XY_m) for i in range(len(XY_m))])
plt.ylabel('XY')
plt.subplot(223)
plt.plot(amplitude(YY_m))
plt.plot([np.mean(YY_m)-n_s*np.std(YY_m) for i in range(len(YY_m))])
plt.plot([np.mean(YY_m)+n_s*np.std(YY_m) for i in range(len(YY_m))])
plt.ylabel('YY')
plt.subplot(224)
plt.plot(amplitude(YX_m))
plt.plot([np.mean(YX_m)-n_s*np.std(YX_m) for i in range(len(YX_m))])
plt.plot([np.mean(YX_m)+n_s*np.std(YX_m) for i in range(len(YX_m))])
plt.ylabel('YX')
fig11.savefig(file_name[:-4]+'1(clean).png')

plt.figure()
plt.subplot(221)
plt.plot(R_XX)
plt.subplot(222)
plt.plot(R_XY)
plt.subplot(223)
plt.plot(R_YY)
plt.subplot(224)
plt.plot(R_YX)

plt.show()


from fbs_runtime.application_context import ApplicationContext
from PyQt5.QtWidgets import (QMainWindow, QMessageBox)
from PyQt5.QtCore import (QThread, pyqtSignal)
from PyQt5.QtGui import (QIcon)
from decimal import Decimal, ROUND_FLOOR
from logger import log
from pprint import pprint
from steam import steam
from math import ceil, floor
from gui import (mainwindow, resources)

import urllib.request
import qdarkstyle
import subprocess
import webbrowser
import speedtest
import win32api
import ctypes
import psutil
import xerox
import time
import yaml
import sys
import os

VERSION = "1.0.0"

launchOptions = [
    "-novid", "-nojoy", "-nosteamcontroller", "-softparticlesdefaultoff",
    "-reuse", "-nohltv"
]

dxlevels = ["80", "81", "90", "92", "95 (default)", "98"]

def is_admin():
  try:
    return ctypes.windll.shell32.IsUserAnAdmin()
  except:
    return False

class SpeedTest(QThread):
  signal = pyqtSignal("PyQt_PyObject")

  def __init__(self):
    QThread.__init__(self)

  def run(self):
    s = speedtest.Speedtest()
    s.get_best_server()
    s.download()
    s.upload()
    s.results.share()

    results_dict = s.results.dict()
    # results_dict = {"download": 0, "upload": 0}
    log.debug(results_dict)
    self.signal.emit(results_dict)

class App(QMainWindow):

  def __init__(self, parent=None):
    super(App, self).__init__(parent)
    self.ui = mainwindow.Ui_MainWindow()
    self.ui.setupUi(self)
    self.setWindowTitle("configuraTed v{v}".format(v=VERSION))
    self.ui.nextButton.clicked.connect(self.onNextClick)
    self.ui.prevButton.clicked.connect(self.onPrevClick)
    self.ui.autoRefreshRate.clicked.connect(self.getRefreshRate)
    self.ui.copyBackupPath.clicked.connect(self.copyBackupPath)
    self.ui.copyConfiguratedPath.clicked.connect(self.copyConfiguratedPath)
    self.ui.toggleRefreshRate.stateChanged.connect(self.onToggleRefreshRate)
    self.ui.refreshConnectionSpeed.clicked.connect(self.getConnectionSpeed)
    self.ui.refreshConnectionSpeed.setIcon(QIcon(":/icons/icons/refresh.png"))
    self.toggleButton(self.ui.nextButton, "show")
    self.toggleButton(self.ui.prevButton, "hide")
    self.setStyleSheet(qdarkstyle.load_stylesheet_pyqt5())

    self.custom = ""
    self.configs = {}
    self.backupPath = ""
    log.debug("loading configs.yml")
    try:
      with open("configs.yml", encoding="utf8") as f:
        self.configs = yaml.safe_load(f)
    except Exception as e:
      msg = "Error while trying to get configs.yml: %s" % e
      log.error(msg, exc_info=True)
      QMessageBox.critical(self, "Error", msg)
    log.debug("configs.yml loaded")

    self.ui.configList.addItem("TF2's Default")

    for id, cfg in self.configs["list"].items():
      log.debug("%s = %s" % (id, cfg))
      self.ui.configList.addItem(cfg["name"])

    self.ui.configList.setCurrentIndex(5)

    for dxl in dxlevels:
      self.ui.dxlevelList.addItem(dxl)

    self.ui.connectionType.addItem("Cable")
    self.ui.connectionType.addItem("Wifi")
    self.ui.connectionType.addItem("Mobile")

    self.ui.dxlevelList.setCurrentIndex(4)
    self.ui.pages.setCurrentIndex(0)
    self.shouldTestConnection = True
    self.getRefreshRate()

  def openBackupPath(self):
    subprocess.Popen("explorer /select, \"%s\"" % self.backupPath)

  def openConfiguratedPath(self):
    subprocess.Popen("explorer /select, \"%s\"" % self.configuratedPath)

  def copyBackupPath(self):
    xerox.copy(self.backupPath)

  def copyConfiguratedPath(self):
    xerox.copy(self.configuratedPath)

  def getRefreshRate(self):
    log.debug("getting refresh rate...")
    try:
      displayfrequency = getattr(
          win32api.EnumDisplaySettings(win32api.EnumDisplayDevices().DeviceName,
                                       -1), 'DisplayFrequency')
      self.ui.refreshRate.setText(str(displayfrequency))
      log.debug("refresh rate: %s" % displayfrequency)
    except Exception as e:
      msg = "Error while trying to get the refresh rate: %s" % e
      log.error(msg, exc_info=True)
      QMessageBox.critical(self, "Error", msg)

  def onToggleRefreshRate(self):
    self.ui.refreshRate.setDisabled(self.ui.toggleRefreshRate.isChecked())

  def clamp(self, actual, mi, ma):
    return max(min(actual, ma), mi)

  def onSpeedTestDone(self, data):
    self.shouldTestConnection = False
    self.ui.statusbar.showMessage("Done, click next to continue.")
    log.debug("done: %s" % data)
    self.toggleButton(self.ui.nextButton, "show")
    self.ui.refreshConnectionSpeed.setEnabled(True)

    d = str(data["download"] / 1000000)
    d = d[0:d.index(".") + 2]

    u = str(data["upload"] / 1000000)
    u = u[0:u.index(".") + 2]
    self.ui.connectionSpeed.setText("{down} Down / {up} Up (Mbps)".format(
        down=d, up=u))

    speed = min(float(d), float(u))

    # rate
    rate = float(speed * 125000)
    rate = self.clamp(rate, 40000, 324192)
    self.custom_net = "\nrate " + str(rate)

    # net split
    net_split = float(speed * 125000)
    net_split = self.clamp(net_split, 3500, 324192)
    self.custom_net += "\nnet_splitpacket_maxrate " + str(net_split)

    # net clear
    self.custom_net += "\nnet_maxcleartime " + str(
        max(4000 / 786432, 4000 / (float(u) * 125000)))
    log.debug(self.custom_net)

  def toggleButton(self, btn, state):
    if state == "hide":
      btn.setStyleSheet(
          "background-color: rgba(0, 0, 0, 0); color: rgba(0, 0, 0, 0);")
      btn.setEnabled(False)
    else:
      btn.setStyleSheet("")
      btn.setEnabled(True)

  def getConnectionSpeed(self):
    try:
      self.st = SpeedTest()
      self.st.signal.connect(self.onSpeedTestDone)
      self.st.start()
      self.ui.statusbar.showMessage("Testing connection speed...")
      self.toggleButton(self.ui.nextButton, "hide")
      self.ui.refreshConnectionSpeed.setEnabled(False)
    except:
      e = sys.exc_info()[0]
      log.error("Error while trying to get the connection speed: %s" % e)
      QMessageBox.critical(self, "Error", "%s" % e)

  def onNextClick(self):
    log.debug("next")
    c = self.ui.pages.currentIndex()
    if c == 0 and self.shouldTestConnection:
      self.getConnectionSpeed()
      return
    elif c == 3:
      self.Install()

    self.ui.pages.setCurrentIndex(c + 1)
    self.toggleButton(self.ui.prevButton, "show")

    if (c + 1) == 3:
      self.ui.nextButton.setText("Install")
    else:
      self.ui.nextButton.setText("Next")

  def onPrevClick(self):
    log.debug("prev")
    c = self.ui.pages.currentIndex()
    if c > 0:
      self.ui.pages.setCurrentIndex(c - 1)
      if c - 1 == 0:
        self.ui.prevButton.setEnabled(False)

      self.ui.nextButton.setText("Next")

  def Install(self):
    log.debug("install")
    self.toggleButton(self.ui.prevButton, "hide")
    self.toggleButton(self.ui.nextButton, "hide")

    buttonReply = QMessageBox.question(
        self, "Confirmation", "Do you wanna proceed with the installation?",
        QMessageBox.Yes | QMessageBox.No, QMessageBox.No)
    if buttonReply == QMessageBox.No:
      return

    self.ui.statusbar.showMessage("Installing...")
    shouldReset = self.ui.resetBox.isChecked()
    shouldLaunch = self.ui.launchBox.isChecked()
    dxlevel = self.ui.dxlevelList.currentText()
    cfg = self.ui.configList.currentText()

    try:
      if shouldReset:
        self.ui.statusbar.showMessage("Reset TF2: finding steam's path...")
        self.errorMessage = "steam.findGame()"
        steam.findGame()

        self.ui.statusbar.showMessage("Reset TF2: backing up user's config...")
        self.errorMessage = "steam.backupUserConfig()"
        self.backupPath = steam.backupUserConfig()

        self.ui.statusbar.showMessage(
            "Reset TF2: disabling steam's cloud for tf2...")
        self.errorMessage = "steam.disableCloud()"
        steam.disableCloud()

        self.ui.statusbar.showMessage("Reset TF2: validating tf2's files...")
        self.errorMessage = "steam.validateGameFiles()"
        steam.validateGameFiles()

        self.ui.statusbar.showMessage(
            "Reset TF2: generating tf2's default config...")
        self.errorMessage = "steam.autoConfig()"
        steam.autoConfig()

        time.sleep(4)

      if shouldLaunch or dxlevel[0:2] != "95":
        self.ui.statusbar.showMessage(
            "Setting tf2's dxlevel to {dxlevel}...".format(dxlevel=dxlevel))
        self.errorMessage = "steam.autoConfig set_dxlevel"
        steam.autoConfig(mode="set_dxlevel", dxlevel=dxlevel[0:2])

        self.ui.statusbar.showMessage(
            "Setting tf2's launch options".format(dxlevel=dxlevel))
        launchOptions.append("-freq " + self.ui.refreshRate.text())
        self.errorMessage = "steam.setLaunchOptions()"
        steam.setLaunchOptions(launchOptions)

      self.ui.statusbar.showMessage("Installing config: {cfg}".format(cfg=cfg))
      self.errorMessage = "installConfig()"
      self.installConfig()

      self.ui.backupPath.setText(self.backupPath)
      self.ui.statusbar.showMessage("Done, you can play tf2 now.")
      self.ui.pages.setCurrentIndex(self.ui.pages.currentIndex() + 1)
    except Exception as e:
      msg = "%s: %s" % (self.errorMessage, e)
      log.error(msg, exc_info=True)
      QMessageBox.critical(self, "Error", msg)

  def installConfig(self):
    log.debug("installing config")
    current = self.ui.configList.currentText()
    if current == "TF2's Default":
      self.configuratedPath = steam.findGame() + "cfg\\"
      return

    cfg = {}
    for _id, c in self.configs["list"].items():
      if c["name"] == current:
        cfg = c
        break

    customPath = steam.findGame() + "custom\\"
    self.configuratedPath = customPath + "configurated\\"

    log.debug("downloading cfg")
    ndir = self.configuratedPath + "cfg\\"
    log.debug("making dirs: %s" % ndir)
    os.makedirs(ndir, exist_ok=True)

    # when type is autoexec it downloads directly to: tf/custom/configurated/
    if cfg["type"] == "autoexec":
      urllib.request.urlretrieve(cfg["url"] + cfg["file"], ndir + cfg["file"])
    else:
      urllib.request.urlretrieve(cfg["url"] + cfg["file"],
                                 customPath + cfg["file"])

    customPath = self.configuratedPath
    self.ui.configuratedPath.setText(customPath)

    if "mastercomfig" in cfg["name"].lower():
      log.debug(
          "cfg is mastercomfig, creating cfg's custom files at tf/custom/configurated/cfg"
      )
      files = [
          "custom", "scout", "soldier", "demoman", "engineer", "heavyweapons",
          "medic", "pyro", "sniper", "spy", "game_overrides", "listenserver"
      ]

      if not ("experimental" in cfg["name"].lower()):
        log.debug("adding _c to cfg file names")
        for i in range(len(files)):
          if files[i] != "custom":
            files[i] = files[i] + "_c"

      for name in files:
        open(customPath + "cfg\\" + name + ".cfg", "w", encoding="utf8").close()

      f = open(customPath + "cfg\\" + "custom.cfg", "w", encoding="utf8")
      f.write(self.custom + self.custom_net)
      f.close()

if __name__ == "__main__":
  appctxt = ApplicationContext()  # 1. Instantiate ApplicationContext
  window = App()
  window.show()
  exit_code = appctxt.app.exec_()  # 2. Invoke appctxt.app.exec_()
  sys.exit(exit_code)


import tushare as ts
import pandas as pd

loc = 'C:\\Users\\Caesar\\Desktop\\Data\\share.xlsx'
df = pd.read_excel(loc, 0)

df_out = pd.DataFrame(columns=['code', 'value'])

for codes in df[['code']].values:
    #print(codes[0])
    #20日历史数据
    code = str(codes[0]).zfill(6)
    df = ts.get_hist_data(code)
    #print(type(df))
    if df is not None :
        turnover_df = df[['turnover']]
        if turnover_df.empty != True :
            turnover = turnover_df.values[0][0]
            if turnover>=15:
                row = {'code': code, 'value': turnover}
                df_out = df_out.append(row, ignore_index=True)
                print(code)
                print(str(turnover))

path = 'C:\\Users\\Caesar\\Desktop\\test.xlsx'
df_out.to_excel(path)


#!/usr/bin/env python
# -*- coding: utf-8 -*-

import json
import time
from pyspark.ml import Pipeline
from pyspark.ml.feature import StringIndexer
from pyspark.sql import SparkSession, DataFrame as SparkDataFrame
from pyspark.sql.functions import mean as _mean, col, sort_array, collect_list, struct, udf
from pyspark.mllib.evaluation import RankingMetrics
from pyspark.sql.types import StringType, ArrayType, StructType, StructField, FloatType, IntegerType
import argparse

class Bias:
    '''
        popularity baseline model.

        The predictor is based on the following algorithm:
        𝑏𝑢𝑖=𝜇+𝑏𝑢+𝑏𝑖

        where 𝜇 is the global mean rating. 𝑏𝑖 is the item bias, and 𝑏𝑢 is the user bias.
        hey are computed as follows:

        𝑏𝑖=∑𝑢∈R(𝑖)(𝑟𝑢𝑖−𝜇) / 𝜆2+|R(𝑖)|.

        𝑏𝑢=∑𝑖∈R(𝑢)(𝑟𝑢𝑖−𝜇−𝑏𝑖) / 𝜆3+|R(𝑢)|.

        Args:
            item: whether to compute item biases
            users: whether to compute user biases
            damping(number or tuple):
                Bayesian damping to apply to computed biases.  Either a number, to
                damp both user and item biases the same amount, or a (user,item) tuple
                providing separate damping values.
        Attributes:
            mean: The global mean rating.
            item_offsets_(DataFrame): The item offsets (𝑏𝑖 values)
            user_offsets_(DataFrame): The user offsets (𝑏𝑢 values)
    '''

    def __init__(self, item=True, users=True, damping=0.0):
        self.item = item
        self.users = users
        if isinstance(damping, tuple):
            self.damping = damping
            self.user_damping, self.item_damping = damping
        else:
            self.damping = damping
            self.user_damping = damping
            self.item_damping = damping

        if self.user_damping < 0:
            raise ValueError("user damping must be non-negative")
        if self.item_damping < 0:
            raise ValueError("item damping must be non-negative")

    def torating(self, data):
        describe = data.describe(['count']).collect()
        stddev = float(describe[2]['count'])
        mean = float(describe[1]['count'])
        max = mean + 2 * stddev
        min = mean - 2 * stddev
        mul = float((max - min) / 100)

        @udf(FloatType())
        def counttrans(count):
            if count < min:
                count = min
            elif count > max:
                count = max

            percentiles = (count - min) / mul
            if percentiles > 100:
                percentiles = 100
            elif percentiles < 0:
                percentiles = 0
            result = percentiles / 100 * 5

            return result if result > 0 else 0.05

        result = data.withColumn('count', counttrans(col('count')))

        return result

    def fit(self, ratings):
        """
        Train the bias model on some rating data.
        Args:
            ratings (DataFrame): the dataframe of song dataset.
        Returns:
            Bias: the fit bias object.
        """
        ratings_stat = ratings.select(_mean(col('count'))).collect()
        self.mean = ratings_stat[0]['avg(count)']
        print(f'global mean: {self.mean}')
        nrate = ratings.withColumn('minus', col('count') - self.mean).select('minus', 'user_id', 'track_id')

        if self.item:
            self.item_offsets_ = self._group(nrate.groupby('track_id'), self.item_damping, 'track_id')
        else:
            self.item_offsets_ = None

        if self.users:
            if self.item_offsets_ is not None:
                nrate \
                    .join(self.item_offsets_, on='track_id', how='inner') \
                    .withColumn('minus', col('minus') - col('track_id_off')) \
                    .select('minus', 'user_id', 'track_id')

                self.user_offsets_ = self._group(nrate.groupby('user_id'), self.user_damping, 'user_id')
        else:
            self.user_offsets_ = None

        return self

    def predict(self, ratings):
        """
        Do predict for the giving item and user.
        The method use the result of fit() and do not recalculate the 𝑏𝑖 and 𝑏𝑢
        Args:
            ratings (DataFrame): the dataframe of song dataset you want to predict,
                            must have track_id column and user_id column.
        Returns:
            rvsp(DataFrame): the DataFrame with predict result, which in prediction column
        """

        if self.item_offsets_ is not None:
            rvsp = self.item_offsets_.join(ratings, on='track_id', how='left') \
                .fillna({'track_id_off': 0}) \
                .withColumn('prediction', self.mean + col('track_id_off'))
        if self.user_offsets_ is not None:
            rvsp = self.user_offsets_.join(rvsp, on='user_id', how='left') \
                .fillna({'user_id_off': 0}) \
                .withColumn('prediction', col('prediction') + col('user_id_off')) \
                .select('prediction', 'user_id', 'track_id', 'count')
        return rvsp

    def recommendForUserSubset(self, data: SparkDataFrame, k: int = 50, colName: str = 'prediction'):

        stringIndexer = StringIndexer(inputCol="track_id", outputCol="track_id_indexed", handleInvalid="skip", stringOrderType="frequencyDesc")
        model = stringIndexer.fit(data)
        data = model.transform(data).withColumn('track_id_indexed', col('track_id_indexed').cast(IntegerType()))

        limit = udf(lambda z: z[0:k], ArrayType(IntegerType(), True))
        if colName == 'count':
            recommend = data.groupBy('user_id').agg(
                sort_array(collect_list(struct(colName, 'track_id_indexed')), asc=False).alias('list')) \
                .withColumn('labels', col('list.track_id_indexed')).drop('list')
        else:
            recommend = data.groupBy('user_id').agg(
                sort_array(collect_list(struct(colName, 'track_id_indexed')), asc=False).alias('list')) \
                .withColumn('labels', col('list.track_id_indexed')).drop('list')

        return recommend

    def score(self, data: SparkDataFrame, k: int = 50) -> dict:
        predictions = self.recommendForUserSubset(self.predict(data), k, 'prediction') \
            .withColumnRenamed('labels', 'prediction')
        label = self.recommendForUserSubset(data, k, 'count')
        results = label.join(predictions, on='user_id', how='left')
        # results.select('prediction', 'labels').show(n=30, truncate=False, vertical=True)
        _array = udf(lambda z: list() if z is None else z, ArrayType(IntegerType(), True))
        result_rdd = results.select('prediction', 'labels').withColumn('labels', _array(col('labels'))).withColumn('prediction', _array(col('prediction'))).rdd
        # print(results.select('prediction', 'labels').rdd.take(5))
        metrics = RankingMetrics(result_rdd)

        ###
        ## NOTE: RankingEvaluator does not work with PySpark 2.4.0.
        ## See https://spark.apache.org/docs/latest/api/python/reference/api/pyspark.ml.evaluation.RankingEvaluator.html

        # evaluator = RankingEvaluator(predictionCol='predictions', labelCol='labels')
        # value = evaluator.evaluate(results,
        #                            {evaluator.metricName: 'precisionAtK', evaluator.k: k})

        return {
            'map': metrics.meanAveragePrecision,
            f'ndcgAt{k}': metrics.ndcgAt(k),
            f'precisionAt{k}': metrics.precisionAt(k)
        }

    def save(self, spark, path):
        data = ([(self.mean,)])

        schema = StructType([StructField("mean", FloatType(), True)])

        df = spark.createDataFrame(data=data, schema=schema)
        df.write.mode('overwrite').parquet(path + 'mean.parquet')
        self.user_offsets_.write.mode('overwrite').parquet(path + 'user_offsets_.parquet')
        self.item_offsets_.write.mode('overwrite').parquet(path + 'item_offsets_.parquet')
        print(f'save {self.mean}, {self.user_offsets_}, {self.item_offsets_} to {path}')

    def load(self, spark, path):
        model = spark.read.parquet(path + 'mean.parquet')
        model_stat = model.select('*').collect()
        self.mean = model_stat[0]['mean']
        self.user_offsets_ = spark.read.parquet(path + 'user_offsets_.parquet')
        self.item_offsets_ = spark.read.parquet(path + 'item_offsets_.parquet')
        print(f'load {self.mean}, {self.user_offsets_}, {self.item_offsets_} from {path}')

        return self

    def _group(self, group, damping, id):
        if damping is not None and damping > 0:
            count = group \
                .count() \
                .withColumn('add', col('count') + self.item_damping)

            sum = group \
                .sum('minus') \
                .withColumnRenamed('sum(minus)', 'sum')

            off = sum \
                .join(count, on=id, how='inner') \
                .withColumn(id + '_off', col('sum') / col('add')) \
                .select(id, id + '_off')

            return off
        else:
            return group.mean('minus')

def main(spark, train_dir, test_dir, val_dir, damping, k, save_dir, load_dir, torating: bool=False):
    """
    test of Class Bias
    """
    test = Bias(item=True, users=True, damping=damping)
    raw_train_data = spark.read.parquet(train_dir)
    train_data = test.torating(raw_train_data) if torating else raw_train_data
    if (load_dir is not None):
        predictions = test.load(spark, load_dir)
    else:
        start = time.time()
        predictions = test.fit(train_data)
        end = time.time()
        print(f'fit function spends {end - start}')
        if (save_dir is not None):
            predictions.save(spark, save_dir)

    results = {}

    train_metrics = predictions.score(train_data, k)
    for m, v in train_metrics.items():
        results[f'train/{m}'] = v

    if val_dir:
        val_metrics = predictions.score(spark.read.parquet(val_dir), k)
        for m, v in val_metrics.items():
            results[f'val/{m}'] = v

    if test_dir:
        test_metrics = predictions.score(spark.read.parquet(test_dir), k)
        for m, v in test_metrics.items():
            results[f'test/{m}'] = v

    print(json.dumps(results, indent=2))

if __name__ == "__main__":
    '''

    test of Class Bias
  use this command to run this script
  ```shell
  spark-submit --conf spark.executor.memory=4g bias.py \
   -tr hdfs:/user/${USER}/quarantini/${TrainFile} \
    -te hdfs:/user/${USER}/quarantini/${TestFile} \
     -du ${USER_DAMPING} -di ${ITEM_DAMPING}
  ```
  '''
    parser = argparse.ArgumentParser(description='popularity baseline model bias.py test')
    parser.add_argument('-tr', '--train', type=str, dest='train_dir', metavar='',
                        default=None, help='Train data directory')
    parser.add_argument('-te', '--test', type=str, dest='test_dir', metavar='',
                        default=None, help='Test data directory')
    parser.add_argument('-v', '--val', type=str, dest='val_dir', metavar='',
                        default=None, help='Val data directory')
    parser.add_argument('-du', '--user_dampling', type=int, dest='user_damping', metavar='',
                        default=None,
                        help='damping of bias model')
    parser.add_argument('-di', '--item_dampling', type=int, dest='item_damping', metavar='',
                        default=None,
                        help='item damping of bias model')
    parser.add_argument('-k', '--k', type=int, dest='k', metavar='',
                        default=None,
                        help='top k recommendation')
    parser.add_argument('-l', '--load', type=str, dest='load_dir', metavar='',
                        default=None,
                        help='load file directory')
    parser.add_argument('-s', '--save', type=str, dest='save_dir', metavar='',
                        default=None,
                        help='save  file directory')
    parser.add_argument('-r', '--rating', type=int, dest='rating', metavar='',
                        default=None,
                        help='transfer to rating or not 1 for ture')
    args = parser.parse_args()

    spark = SparkSession.builder.appName('bias_test').config('spark.blacklist.enabled', False).getOrCreate()

    damping = (args.user_damping, args.item_damping)

    transferrate = True if args.rating == 1 else False

    main(spark=spark, train_dir=args.train_dir, test_dir=args.test_dir, val_dir=args.val_dir, damping=damping, k=args.k,
         save_dir=args.save_dir, load_dir=args.load_dir, torating=transferrate)


import sys
import tornado
import asyncio
import logging
from tornado.ioloop import PeriodicCallback
from tornado.options import define, options
import tornado.platform.asyncio

from .application import BrokerApplication
from .coap import coap_server_init, _check_dead_nodes
from .logger import logger

logging.basicConfig(level=logging.DEBUG,
                    format='%(asctime)s - %(name)14s - '
                           '%(levelname)5s - %(message)s')

def parse_command_line():
    """Parse command line arguments for IoT broker application."""
    define("port", default=8000,
           help="Broker port")
    define("max_time", default=120,
           help="Retention time for lost nodes (s).")
    define("debug", default=False,
           help="Enable debug mode.")
    options.parse_command_line()

    if options.debug:
        logger.setLevel(logging.DEBUG)

def run():
    """Start a broker instance."""
    parse_command_line()
    try:
        # Tornado ioloop initialization
        ioloop = asyncio.get_event_loop()
        tornado.platform.asyncio.AsyncIOMainLoop().install()
        PeriodicCallback(_check_dead_nodes, 1000).start()

        # Initialize Coap server
        coap_server_init(options.max_time)

        # Start tornado application
        app = BrokerApplication(options=options)
        app.listen(options.port)
        ioloop.run_forever()
    except KeyboardInterrupt:
        print("Exiting")
        ioloop.stop()
        sys.exit()

if __name__ == '__main__':
    run()


#!/usr/bin/env python

'''
This script takes output from a pairwise (for now) RBBH run, extracts those orthologs from the respective
Trinity.fasta output file and renames transcripts in the second fasta file to the corresponding ones in the first

usage: extract_rename_cRBH_orthologs.py <RBBH_output.txt> <first_fasta_input> <second_fasta_input> <first_fasta_output> <second_fasta_output>

'''
from Bio import SeqIO
import sys
import re
from random import randint

outfile_first = open(sys.argv[4], "w")
outfile_second = open(sys.argv[5], "w")
orthos = open(sys.argv[1], "rU")
first_fasta = open(sys.argv[2], "rU")
second_fasta = open(sys.argv[3], "rU")

genes_first = []
genes_second = []
ortho_second = {}

for line in orthos:
    info = line.split('\t')
    gene_first = info[0]
    gene_second = info[1]
    genes_first.append(gene_first)
    genes_second.append(gene_second)
    ortho_second[gene_second] = gene_first

for record in SeqIO.parse(first_fasta, "fasta"):
    if record.id in genes_first:
        print ('Writing: ' + record.id)
        SeqIO.write(record, outfile_first, "fasta")

for record in SeqIO.parse(second_fasta, "fasta"):
    if record.id in genes_second:
        record.id = ortho_second[record.id]
        print ('Writing: ' + record.id)
        SeqIO.write(record, outfile_second, "fasta")

outfile_first.close()
outfile_second.close()
orthos.close()
first_fasta.close()
second_fasta.close()


"""Module for interacting with MATL, and it's source code."""
import pathlib
from typing import Optional

from matl_online.octave import OctaveSession, OutputCallback
from matl_online.octave import string as octave_string
from matl_online.types import MATLTaskParameters

from .source import get_matl_folder

def matl(
    octave: OctaveSession,
    matl_params: MATLTaskParameters,
    directory: pathlib.Path,
    line_handler: Optional[OutputCallback] = None,
) -> None:
    """Open a session with Octave and manages input/output as well as errors."""

    # Add the folder for the appropriate MATL version
    matl_folder = get_matl_folder(matl_params.version)

    # Ensure the matl folder exists
    assert matl_folder, "MATL folder does not exist"

    # Change directories to the temporary folder so that all temporary
    # files are placed in here and won't interfere with other requests
    with octave.current_directory(directory):
        with octave.paths(matl_folder):
            # Convert the code to a cell array element-per-line
            code = f"{{{','.join([octave_string(x) for x in matl_params.code_lines])}}}"

            octave.run(
                "matl_runner",
                octave_string(matl_params.flags),
                code,
                *[octave_string(x) for x in matl_params.input_lines],
                line_handler=line_handler,
            )


#!/usr/bin/env python3

import asyncio
import random
import sys

resource = ['*', '$', '#']
resource_lock = asyncio.Lock()
#排队锁，不论执行读操作还是写操作都要先获得该锁，在获得资源锁后释放。
queue_lock = asyncio.Lock()
read_count = 0

async def writer(data, name):
    global resource, resource_lock, queue_lock
    #申请获得排队锁。
    await queue_lock.acquire()
    await resource_lock.acquire()
    #释放排队锁。
    queue_lock.release()
    for i in range(3):
        t = random.random()
        await asyncio.sleep(t)
        resource[i] = data[i]
        print(f"{name}: write {data[i]}, now {resource[0] + resource[1] + resource[2]}")
    resource_lock.release()

async def reader(name):
    global resource, resource_lock, read_count, queue_lock
    #申请获得排队锁。
    await queue_lock.acquire()
    read_count = read_count + 1
    if read_count == 1:
        await resource_lock.acquire()
    #释放排队锁。
    queue_lock.release()
    for i in range(3):
        t = random.random()
        await asyncio.sleep(t)
        print(f"{name}: read {resource[i]}")
    read_count = read_count - 1
    if read_count == 0:
        resource_lock.release()

async def main():
    r1 = asyncio.create_task(reader("r1"))
    w1 = asyncio.create_task(writer("abc", "w1"))
    r2 = asyncio.create_task(reader("r2"))
    r3 = asyncio.create_task(reader("r3"))
    w2 = asyncio.create_task(writer("def", "w2"))
    r4 = asyncio.create_task(reader("r4"))
    await asyncio.gather(r1, w1, r2, r3, w2, r4)
    sys.exit(0)

if __name__ == '__main__':
    asyncio.run(main())


import numpy as np
import nltk
from itertools import chain
from nltk.corpus import wordnet

class Wordnet:
    """
    A set of functions used to augment data.

    Typical usage: ::
        >>> import nltk
        >>> nltk.download('punkt')
        >>> nltk.download('wordnet')
        >>> nltk.download('averaged_perceptron_tagger')
        >>> from textaugment import Wordnet
        >>> t = Wordnet(v=True,n=True,p=0.5)
        >>> t.augment('I love school')
        i adore school
    """

    def __init__(self, **kwargs):
        """
        A method to initialize parameters

        :type random_state: int
        :param random_state: seed
        :type v: bool
        :param v: Verb, default is True
        :type n: bool
        :param n: Noun
        :type runs: int
        :param runs: Number of repetition on single text
        :type p: float, optional
        :param p: The probability of success of an individual trial. (0.1<p<1.0), default is 0.5
        :rtype:   None
        :return:  Constructer do not return.
        """

        # Set random state
        if 'random_state' in kwargs:
            self.random_state = kwargs['random_state']
            if isinstance(self.random_state, int):
                np.random.seed(self.random_state)
            else:
                raise TypeError("random_state must have type int, float, str, bytes, or bytearray")

        # Set verb to be default if no values given
        try:
            if "v" not in kwargs and "n" not in kwargs:
                kwargs['v'] = True
                kwargs['n'] = False
            elif "v" in kwargs and "n" not in kwargs:
                kwargs['v'] = True
                kwargs['n'] = False
            elif "v" not in kwargs and "n" in kwargs:
                kwargs['n'] = True
                kwargs['v'] = False
            if "runs" not in kwargs:
                kwargs['runs']=1

        except KeyError:
            raise

        try:
            if "p" in kwargs:
                if type(kwargs['p']) is not float:
                    raise TypeError("p represent probability of success and must be a float from 0.1 to 0.9. E.g p=0.5")
                elif type(kwargs['p']) is float:
                    self.p = kwargs['p']
            else:
                kwargs['p'] = 0.5  # Set default value
        except KeyError:
            raise

        self.p = kwargs['p']
        self.v = kwargs['v']
        self.n = kwargs['n']
        self.runs = kwargs['runs']

    def geometric(self, data):
        """
        Used to generate Geometric distribution.

        :type data: list
        :param data: Input data
        :rtype:   ndarray or scalar
        :return:  Drawn samples from the parameterized Geometric distribution.
        """

        data = np.array(data)
        first_trial = np.random.geometric(p=self.p, size=data.shape[0]) == 1  # Capture success after first trial
        return data[first_trial]

    def replace(self, data, lang):
        """
        The method to replace words with synonyms

        :type data: str
        :param data: sentence used for data augmentation
        :rtype:   str
        :return:  The augmented data
        :type lang: str
        :param lang: choose lang
        """
        data = data.lower().split()
        data_tokens = [[i, x, y] for i, (x, y) in enumerate(nltk.pos_tag(data))]  # Convert tuple to list
        if self.v:
            for loop in range(self.runs):
                words = [[i, x] for i, x, y in data_tokens if y[0] == 'V']
                words = [i for i in self.geometric(data=words)]  # List of selected words
                if len(words) >= 1:  # There are synonyms
                    for word in words:
                        synonyms1 = wordnet.synsets(word[1], wordnet.VERB, lang=lang)  # Return verbs only
                        synonyms = list(set(chain.from_iterable([syn.lemma_names(lang=lang) for syn in synonyms1])))
                        synonyms_ = []  # Synonyms with no underscores goes here
                        for w in synonyms:
                            if '_' not in w:
                                synonyms_.append(w)  # Remove words with underscores
                        if len(synonyms_) >= 1:
                            synonym = self.geometric(data=synonyms_).tolist()
                            if synonym:  # There is a synonym
                                data[int(word[0])] = synonym[0].lower()  # Take the first success

        if self.n:
            for loop in range(self.runs):
                words = [[i, x] for i, x, y in data_tokens if y[0] == 'N']
                words = [i for i in self.geometric(data=words)]  # List of selected words
                if len(words) >= 1:  # There are synonyms
                    for word in words:
                        synonyms1 = wordnet.synsets(word[1], wordnet.NOUN, lang=lang)  # Return nouns only
                        synonyms = list(set(chain.from_iterable([syn.lemma_names(lang=lang) for syn in synonyms1])))
                        synonyms_ = []  # Synonyms with no underscores goes here
                        for w in synonyms:
                            if '_' not in w:
                                synonyms_.append(w)  # Remove words with underscores
                        if len(synonyms_) >= 1:
                            synonym = self.geometric(data=synonyms_).tolist()
                            if synonym:  # There is a synonym
                                data[int(word[0])] = synonym[0].lower()  # Take the first success

        return " ".join(data)

    def augment(self, data, lang="eng"):
        """
        Data augmentation for text. Generate new dataset based on verb/nouns synonyms.

        :type data: str
        :param data: sentence used for data augmentation
        :rtype:   str
        :return:  The augmented data
        :type lang: str
        :param lang: choose lang
        """
        # Error handling
        if type(data) is not str:
            raise TypeError("Only strings are supported")
        data = self.replace(data, lang)
        return data


import os
import unittest
import json
from flask_sqlalchemy import SQLAlchemy
from flaskr import create_app
from models import setup_db, Question, Category

class TriviaTestCase(unittest.TestCase):
    """This class represents the trivia test case"""

    def setUp(self):
        """Define test variables and initialize app."""
        self.app = create_app()
        self.client = self.app.test_client
        self.database_host = '127.0.0.1:5432'
        self.database_username = 'postgres'
        self.database_name = 'trivia_test'
        self.database_path = 'postgresql+psycopg2://{}@{}/{}'.format(
            self.database_username,
            self.database_host,
            self.database_name)
        setup_db(self.app, self.database_path)

        self.new_question = {
            "question": "Which two are names of Santa's reindeer?",
            "answer": "Comet and Cupid",
            "category": "3",
            "difficulty": 1
        }

        # binds the app to the current context
        with self.app.app_context():
            self.db = SQLAlchemy()
            self.db.init_app(self.app)
            # create all tables
            self.db.create_all()

    def tearDown(self):
        pass

    def test_get_paginated_questions(self):
        res = self.client().get('/questions')
        data = json.loads(res.data)

        self.assertEqual(res.status_code, 200)
        self.assertEqual(data['success'], True)
        self.assertTrue(data['total_questions'])
        self.assertTrue(len(data['questions']))

    def test_get_paginated_categories(self):
        res = self.client().get('/categories')
        data = json.loads(res.data)

        self.assertEqual(res.status_code, 200)
        self.assertEqual(data['success'], True)
        self.assertTrue(data['total_categories'])
        self.assertTrue(len(data['categories']))

# test with wrong values
    def test_404_sent_requesting_beyond_valid_page(self):
        res = self.client().get('/questions?page=1000')
        data = json.loads(res.data)

        self.assertEqual(res.status_code, 404)
        self.assertEqual(data['success'], False)
        self.assertEqual(data['message'], 'resource not found')
# test search with and without results

    def test_get_question_search_with_results(self):
        res = self.client().post('/questions', json={'searchTerm': 'largest'})
        data = json.loads(res.data)

        self.assertEqual(res.status_code, 200)
        self.assertEqual(data['success'], True)
        self.assertTrue(data['total_questions'])
        self.assertEqual(len(data['questions']), 1)

    def test_get_question_search_withoutresults(self):
        res = self.client().post('/questions', json={'searchTerm': 'jdhfkf'})
        data = json.loads(res.data)
        self.assertEqual(res.status_code, 200)
        self.assertEqual(data['total_questions'], 0)
        self.assertEqual(data['success'], True)
# test to new question and error

    def test_create_new_question(self):
        res = self.client().post('/questions', json=self.new_question)
        data = json.loads(res.data)
        pass

    def test_422_if_question_creation_fails(self):
        res = self.client().post('/questions', json=self.new_question)
        data = json.loads(res.data)
        pass

# test with questions by category and error
    def test_questions_by_category(self):
        res = self.client().get('/categories/4/questions')
        data = json.loads(res.data)

        self.assertEqual(data['success'], True)
        self.assertEqual(data['total_questions'], 4)
        self.assertEqual(data['current_category'], 4)

    def test_404_questions_by_category(self):
        res = self.client().get('/categories/1000/questions')
        data = json.loads(res.data)

        self.assertEqual(data['error'], 404)
        self.assertEqual(data['message'], 'resource not found')
        self.assertEqual(data['success'], False)
# test for quizz and error

    def test_quiz_questions(self):
        play_quiz = {
            'previous_questions': [],
            'quiz_category': {
                'type': 'sports',
                'id': '6'}}
        res = self.client().post('/play', json=play_quiz)
        data = json.loads(res.data)

        self.assertEqual(res.status_code, 200)
        self.assertEqual(data['success'], True)

    def test_422_quiz_questions(self):
        res = self.client().post('/play', json={})
        data = json.loads(res.data)

        self.assertEqual(res.status_code, 422)
        self.assertEqual(data['success'], False)
        self.assertEqual(data['message'], 'unprocessable')

# test to delete with correct value and does not exist value
    def test_delete_question(self):
        question = Question(question='test_question', answer='test_answer',
                            difficulty=2, category=3)
        question.insert()
        question_id = question.id

        res = self.client().delete(f'/questions/{question_id}')
        data = json.loads(res.data)

        question = Question.query.filter(
            Question.id == question.id).one_or_none()

        self.assertEqual(res.status_code, 200)
        self.assertEqual(data['success'], True)
        self.assertEqual(data['deleted'], question_id)
        self.assertEqual(question, None)

    def test_404_if_question_does_not_exist(self):
        res = self.client().delete('/questions/5000')
        data = json.loads(res.data)

        self.assertEqual(res.status_code, 404)
        self.assertEqual(data['success'], False)
        self.assertEqual(data['message'], 'resource not found')

# Make the tests conveniently executable
if __name__ == "__main__":
    unittest.main()


import pygame
import sys
import random
import os
os.chdir("C:\\Users\\박예은\\AppData\\Local\\Programs\\Python\\Python37\\One_Big_Archive\\images")

class MyBallClass(pygame.sprite.Sprite):
    def __init__(self, image_file, location, speed):
        pygame.sprite.Sprite.__init__(self)     # call Sprite initializer ; 스프라이트 초기화
        self.image = pygame.image.load(image_file)
        self.rect = self.image.get_rect()
        self.rect.left, self.rect.top = location
        self.speed = speed

    def move(self):
        self.rect = self.rect.move(self.speed)
        if self.rect.left < 0 or self.rect.left > width:
            self.speed[0] = -self.speed[0]

        if self.rect.top < 0 or self.rect.bottom > height:
            self.speed[1] = -self.speed[1]

size = width, height = 640, 480
screen = pygame.display.set_mode(size)
# screen = pygame.display.set_mode([640, 480])
# print(width, type(width))
# print(size, type(size))
screen.fill([255, 255, 255])
img_file = "beach_ball.png"
balls = []
for row in range(0, 3):
    for col in range(0, 3):
        location = [col*180+10, row*180+10]
        speed = [random.choice([-2, 2]), random.choice([-2, 2])]
        ball = MyBallClass(img_file, location, speed)
        balls.append(ball)

while True:
    for i in pygame.event.get():
        if i.type == pygame.QUIT:
            sys.exit()
    pygame.time.delay(20)
    screen.fill([255, 255, 255])
    for ball in balls:
        ball.move()
        screen.blit(ball.image, ball.rect)
    pygame.display.flip()