import os
import time
import logging
import random
from tqdm import tqdm, trange
import numpy as np

import torch
import torch.nn.functional as F
from torch.utils.data import (TensorDataset, DataLoader,
                              RandomSampler, SequentialSampler)
from torch.utils.data.distributed import DistributedSampler
try:
    from torch.utils.tensorboard import SummaryWriter
except:
    from tensorboardX import SummaryWriter

from transformers import AdamW, get_linear_schedule_with_warmup
from transformers import glue_compute_metrics as compute_metrics
from transformers.selective_metrics import rcc_auc, rpp
from transformers import glue_convert_examples_to_features as convert_examples_to_features
from transformers import glue_output_modes as output_modes
from transformers import glue_processors as processors

from torchprofile.profile import profile_macs

logger = logging.getLogger(__name__)

def load_and_cache_examples(args, task, tokenizer, evaluate=False):
    if args.local_rank not in [-1, 0] and not evaluate:
        torch.distributed.barrier()  # Make sure only the first process in distributed training process the dataset, and the others will use the cache

    if task in ['mnli-mm', 'bmnli-mm']:
        data_dir = args.data_dir[:-3]
    else:
        data_dir = args.data_dir
    processor = processors[task]()
    output_mode = output_modes[task]
    # Load data features from cache or dataset file
    cached_features_file = os.path.join(data_dir, 'cached_{}_{}_{}_{}'.format(
        'dev' if evaluate else 'train',
        list(filter(None, args.model_name_or_path.split('/'))).pop(),
        str(args.max_seq_length),
        str(task)))
    if os.path.exists(cached_features_file) and not args.overwrite_cache:
        logger.info("Loading features from cached file %s", cached_features_file)
        features = torch.load(cached_features_file)
    else:
        logger.info("Creating features from dataset file at %s", data_dir)
        label_list = processor.get_labels()
        if task in ['mnli', 'mnli-mm'] and args.model_type in ['roberta']:
            # HACK(label indices are swapped in RoBERTa pretrained model)
            label_list[1], label_list[2] = label_list[2], label_list[1]
        examples = processor.get_dev_examples(data_dir) if evaluate else processor.get_train_examples(data_dir)
        features = convert_examples_to_features(examples,
                                                tokenizer,
                                                label_list=label_list,
                                                max_length=args.max_seq_length,
                                                output_mode=output_mode,
                                                pad_on_left=bool(args.model_type in ['xlnet']),                 # pad on the left for xlnet
                                                pad_token=tokenizer.convert_tokens_to_ids([tokenizer.pad_token])[0],
                                                pad_token_segment_id=4 if args.model_type in ['xlnet'] else 0,
        )
        if args.local_rank in [-1, 0]:
            logger.info("Saving features into cached file %s", cached_features_file)
            torch.save(features, cached_features_file)

    if args.local_rank == 0 and not evaluate:
        torch.distributed.barrier()  # Make sure only the first process in distributed training process the dataset, and the others will use the cache

    # Convert to Tensors and build dataset
    if (args.train_percentage<100) and (not evaluate):
        features = features[:int(len(features) * args.train_percentage / 100)]
    all_input_ids = torch.tensor([f.input_ids for f in features], dtype=torch.long)
    all_attention_mask = torch.tensor([f.attention_mask for f in features], dtype=torch.long)
    all_token_type_ids = torch.tensor([f.token_type_ids for f in features], dtype=torch.long)
    if output_mode in ["classification", "selective-classification"]:
        all_labels = torch.tensor([f.label for f in features], dtype=torch.long)
    elif output_mode == "regression":
        all_labels = torch.tensor([f.label for f in features], dtype=torch.float)
    all_indices = torch.tensor(range(len(features)), dtype=torch.long)

    dataset = TensorDataset(all_input_ids, all_attention_mask, all_token_type_ids, all_labels, all_indices)
    return dataset

def set_seed(args):
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    if args.n_gpu > 0:
        torch.cuda.manual_seed_all(args.seed)

def get_wanted_result(result):
    if "spearmanr" in result:
        print_result = result["spearmanr"]
    elif "f1" in result:
        print_result = result["f1"]
    elif "mcc" in result:
        print_result = result["mcc"]
    elif "acc" in result:
        print_result = result["acc"]
    else:
        print(result)
        exit(1)
    return print_result

def get_dataset_acc(args, dataset, model):
    # routine doesn't matter since it only affects the loss
    eval_batch_size = args.per_gpu_eval_batch_size * max(1, args.n_gpu)
    eval_sampler = SequentialSampler(dataset) if args.local_rank == -1 else DistributedSampler(dataset)
    eval_dataloader = DataLoader(dataset, sampler=eval_sampler, batch_size=eval_batch_size)

    logits = []
    out_label_ids = []
    for batch in eval_dataloader:
        model.eval()
        batch = tuple(t.to(args.device) for t in batch)
        with torch.no_grad():
            inputs = {'input_ids': batch[0],
                      'attention_mask': batch[1],
                      'labels': batch[3]}
            if args.model_type != 'distilbert':
                inputs['token_type_ids'] = batch[2] if args.model_type in ['bert',
                                                                           'xlnet'] else None  # XLM, DistilBERT and RoBERTa don't use segment_ids
            outputs = model(**inputs)
            tmp_eval_loss, batch_logits = outputs[:3]
        logits.append(batch_logits.detach().cpu().numpy())
        out_label_ids.append(inputs['labels'].detach().cpu().numpy())

    logits = np.concatenate(logits, axis=0)
    out_label_ids = np.concatenate(out_label_ids, axis=0)
    if args.output_mode in ["classification", "selective-classification"]:
        preds = np.argmax(logits, axis=1)
    elif args.output_mode == "regression":
        preds = np.squeeze(logits)

    return preds == out_label_ids  # only works for classification

def train(args, train_dataset, model, tokenizer, routine='raw'):
    """ Train the model """
    if args.local_rank in [-1, 0]:
        tb_writer = SummaryWriter()

    args.train_batch_size = args.per_gpu_train_batch_size * max(1, args.n_gpu)
    train_sampler = RandomSampler(train_dataset) if args.local_rank == -1 else DistributedSampler(train_dataset)
    train_dataloader = DataLoader(train_dataset, sampler=train_sampler, batch_size=args.train_batch_size)

    if args.max_steps > 0:
        t_total = args.max_steps
        args.num_train_epochs = args.max_steps // (len(train_dataloader) // args.gradient_accumulation_steps) + 1
    else:
        t_total = len(train_dataloader) // args.gradient_accumulation_steps * args.num_train_epochs

    # Prepare optimizer and schedule (linear warmup and decay)
    no_decay = ['bias', 'LayerNorm.weight']
    optimizer_grouped_parameters = [
        {'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': args.weight_decay},
        {'params': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
        ]
    optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, eps=args.adam_epsilon)
    scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=args.warmup_steps, num_training_steps=t_total)
    if args.fp16:
        try:
            from apex import amp
        except ImportError:
            raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use fp16 training.")
        model, optimizer = amp.initialize(model, optimizer, opt_level=args.fp16_opt_level)

    # multi-gpu training (should be after apex fp16 initialization)
    if args.n_gpu > 1:
        model = torch.nn.DataParallel(model)

    # Distributed training (should be after apex fp16 initialization)
    if args.local_rank != -1:
        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank],
                                                          output_device=args.local_rank,
                                                          find_unused_parameters=True)

    # Train!
    logger.info("***** Running training *****")
    logger.info("  Num examples = %d", len(train_dataset))
    logger.info("  Num Epochs = %d", args.num_train_epochs)
    logger.info("  Instantaneous batch size per GPU = %d", args.per_gpu_train_batch_size)
    logger.info("  Total train batch size (w. parallel, distributed & accumulation) = %d",
                   args.train_batch_size * args.gradient_accumulation_steps * (torch.distributed.get_world_size() if args.local_rank != -1 else 1))
    logger.info("  Gradient Accumulation steps = %d", args.gradient_accumulation_steps)
    logger.info("  Total optimization steps = %d", t_total)

    if routine.endswith('hist'):
        correctness_record = [[0, 0] for _ in range(len(train_dataset))]
    global_step = 0
    tr_loss, logging_loss = 0.0, 0.0
    model.zero_grad()
    train_iterator = trange(int(args.num_train_epochs), desc="Epoch", disable=args.local_rank not in [-1, 0])
    set_seed(args)  # Added here for reproductibility (even between python 2 and 3)
    for i_epoch in train_iterator:
        epoch_iterator = tqdm(train_dataloader, desc="Iteration", disable=args.local_rank not in [-1, 0])
        for step, batch in enumerate(epoch_iterator):
            if (
                    routine.endswith('hist')
                    and step % (len(epoch_iterator) // 10) == 0
                    and step / len(epoch_iterator) <= 0.9
            ):
                # update correctness record
                new_record = get_dataset_acc(args, train_dataset, model)
                for i in range(len(correctness_record)):
                    correctness_record[i][0] += 1  # denominator
                    correctness_record[i][1] += new_record[i]  # numerator

            model.train()
            batch = tuple(t.to(args.device) for t in batch)
            inputs = {
                'input_ids': batch[0],
                'attention_mask': batch[1],
                'labels': batch[3],
                'routine': routine,
                'i_epoch': i_epoch,
                'step': step,
            }
            if routine.endswith('hist'):
                inputs['history_record'] = torch.tensor([
                    correctness_record[ind][1] / correctness_record[ind][0] for ind in batch[-1]
                ], device=args.device)
            if args.model_type != 'distilbert':
                inputs['token_type_ids'] = batch[2] if args.model_type in ['bert', 'xlnet'] else None  # XLM, DistilBERT and RoBERTa don't use segment_ids
            outputs = model(**inputs)
            loss = outputs[0]  # model outputs are always tuple in transformers (see doc)

            if args.n_gpu > 1:
                loss = loss.mean() # mean() to average on multi-gpu parallel training
            if args.gradient_accumulation_steps > 1:
                loss = loss / args.gradient_accumulation_steps

            if args.fp16:
                with amp.scale_loss(loss, optimizer) as scaled_loss:
                    scaled_loss.backward()
            else:
                loss.backward()

            tr_loss += loss.item()
            if (step + 1) % args.gradient_accumulation_steps == 0:
                if args.fp16:
                    torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), args.max_grad_norm)
                else:
                    torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)

                optimizer.step()
                scheduler.step()  # Update learning rate schedule
                model.zero_grad()
                global_step += 1

            if args.max_steps > 0 and global_step > args.max_steps:
                epoch_iterator.close()
                break
        if args.max_steps > 0 and global_step > args.max_steps:
            train_iterator.close()
            break

    if args.local_rank in [-1, 0]:
        tb_writer.close()

    return global_step, tr_loss / global_step

def mc_dropout(args, eval_dataloader, model, runs=10):
    # the model originally has a dropout rate of 0.1, but set to 0.01 for mc_dropout
    logits = []
    for mc_index in range(runs):
        run_logits = []
        torch.manual_seed(42 * (1+mc_index))  # for different dropout seeds
        for batch in tqdm(eval_dataloader, desc="Evaluating"):
            model.train()  # use train() delibrately so dropout is applied
            batch = tuple(t.to(args.device) for t in batch)

            with torch.no_grad():
                inputs = {
                    'input_ids': batch[0],
                    'attention_mask': batch[1],
                    'labels': batch[3],
                }
                if args.model_type != 'distilbert':
                    inputs['token_type_ids'] = batch[2] if args.model_type in ['bert',
                                                                               'xlnet'] else None  # XLM, DistilBERT and RoBERTa don't use segment_ids
                outputs = model(**inputs)
                _, batch_logits = outputs[:3]
            run_logits.append(batch_logits.detach().cpu().numpy())
        logits.append(np.concatenate(run_logits, axis=0))
    raw_logits = np.array(logits)
    max_logits = np.max(raw_logits, axis=2).transpose(1, 0)
    return raw_logits, max_logits

def evaluate(args, model, tokenizer, prefix="", routine='raw'):
    # Loop to handle MNLI double evaluation (matched, mis-matched)

    eval_task = args.task_name
    eval_output_dir = args.output_dir + ("-MM" if "mm" in args.task_name else "")

    results = {}
    eval_dataset = load_and_cache_examples(args, eval_task, tokenizer, evaluate=True)

    if not os.path.exists(eval_output_dir) and args.local_rank in [-1, 0]:
        os.makedirs(eval_output_dir)

    args.eval_batch_size = args.per_gpu_eval_batch_size * max(1, args.n_gpu)
    # Note that DistributedSampler samples randomly
    eval_sampler = SequentialSampler(eval_dataset) if args.local_rank == -1 else DistributedSampler(eval_dataset)
    eval_dataloader = DataLoader(eval_dataset, sampler=eval_sampler, batch_size=args.eval_batch_size)

    # multi-gpu eval
    if args.n_gpu > 1:
        model = torch.nn.DataParallel(model)

    # Eval!
    logger.info("***** Running evaluation {} *****".format(prefix))
    logger.info("  Num examples = %d", len(eval_dataset))
    logger.info("  Batch size = %d", args.eval_batch_size)
    eval_loss = 0.0
    nb_eval_steps = 0
    logits = []
    probs = []
    out_label_ids = []
    start_time = time.time()
    for batch in tqdm(eval_dataloader, desc="Evaluating"):
        model.eval()
        batch = tuple(t.to(args.device) for t in batch)

        with torch.no_grad():
            inputs = {
                'input_ids': batch[0],
                'attention_mask': batch[1],
                'labels': batch[3],
                'routine': routine,
            }
            if args.model_type != 'distilbert':
                inputs['token_type_ids'] = batch[2] if args.model_type in ['bert', 'xlnet'] else None  # XLM, DistilBERT and RoBERTa don't use segment_ids
            outputs = model(**inputs)
            tmp_eval_loss, batch_logits = outputs[:3]
            batch_probs = F.softmax(batch_logits, dim=1)
            eval_loss += tmp_eval_loss.mean().item()
        nb_eval_steps += 1
        logits.append(batch_logits.detach().cpu().numpy())
        probs.append(batch_probs.detach().cpu().numpy())
        out_label_ids.append(inputs['labels'].detach().cpu().numpy())

    eval_time = time.time() - start_time
    eval_loss = eval_loss / nb_eval_steps
    logits = np.concatenate(logits, axis=0)
    probs = np.concatenate(probs, axis=0)
    out_label_ids = np.concatenate(out_label_ids, axis=0)

    if args.output_mode == "selective-classification":
        th_results = {}  # threshold -> score
        for threshold in np.arange(0.7, 1.0, 0.001):
            preds = np.argmax(logits, axis=1)
            for i in range(len(preds)):
                if max(probs[i]) < threshold:
                    preds[i] = model.num_labels - 1
            th_results[threshold] = compute_metrics(eval_task, preds, out_label_ids)['acc']

        best_th, best_acc = 0, 0
        for th, acc in th_results.items():
            # print(th, acc)
            if acc > best_acc:
                best_th, best_acc = th, acc
        print({
            'best_th': best_th,
            'best_acc': best_acc,
        })

    if args.output_mode in ["classification", "selective-classification"]:
        preds = np.argmax(logits, axis=1)
    elif args.output_mode == "regression":
        preds = np.squeeze(logits)
    result = compute_metrics(eval_task, preds, out_label_ids)
    print_result = get_wanted_result(result)

    conf = probs.max(axis=1)
    risk_binary = (preds != out_label_ids).astype(int)
    rcc_auc_value = rcc_auc(conf, risk_binary, 'SR', args)
    rpp_value = rpp(conf, risk_binary)

    print({
        "eval_time": eval_time,
        "result": print_result,
        "rcc_auc": rcc_auc_value,
        "rpp": rpp_value,
    })
    results.update(result)

    if args.top2_diff:
        sorted_probs = np.sort(probs, axis=1)
        t2d_conf = sorted_probs[:, -1] - sorted_probs[:, -2]
        t2d_rcc_auc_value = rcc_auc(t2d_conf, risk_binary, None, args)
        t2d_rpp_value = rpp(t2d_conf, risk_binary)
        print({
            't2d_rcc_auc': t2d_rcc_auc_value,
            't2d_rpp': t2d_rpp_value,
        })

    if args.mc_dropout:
        _, mc_max_logits = mc_dropout(args, eval_dataloader, model)
        mc_conf = -np.var(mc_max_logits, axis=1)
        mc_rcc_auc_value = rcc_auc(mc_conf, risk_binary, 'MC', args)
        mc_rpp_value = rpp(mc_conf, risk_binary)
        print({
            'mc_rcc_auc': mc_rcc_auc_value,
            'mc_rpp': mc_rpp_value,
        })

    if args.multi_mc_dropout:
        save_fname = build_file(
            args.plot_data_dir,
            'saved_data',
            os.path.basename(args.data_dir),
            args.model_name_or_path + '-mc_dropout.npy'
        )
        record_into_file(save_fname, 'sr', rcc_auc_value)

        result_dict = {}
        runs = [2, 5, 10, 20, 50, 100]
        for run in runs:
            _, mc_max_logits = mc_dropout(args, eval_dataloader, model, run)
            mc_conf = -np.var(mc_max_logits, axis=1)
            mc_rcc_auc_value = rcc_auc(mc_conf, risk_binary)
            # mc_rpp_value = rpp(mc_conf, risk_binary)
            result_dict[run] = mc_rcc_auc_value
            print(run, mc_rcc_auc_value)
        record_into_file(save_fname, args.dropout_prob, result_dict)

    output_eval_file = os.path.join(eval_output_dir, prefix, "eval_results.txt")
    with open(output_eval_file, "w") as writer:
        logger.info("***** Eval results {} *****".format(prefix))
        for key in sorted(result.keys()):
            logger.info("  %s = %s", key, str(result[key]))
            writer.write("%s = %s\n" % (key, str(result[key])))

    return results

def multi_stage_evaluate(args, model, super_model, tokenizer,
                         prefix="", profile=True):
    # Loop to handle MNLI double evaluation (matched, mis-matched)

    eval_task = args.task_name
    eval_output_dir = args.output_dir + ("-MM" if "mm" in args.task_name else "")

    results = {}
    eval_dataset = load_and_cache_examples(args, eval_task, tokenizer, evaluate=True)

    if not os.path.exists(eval_output_dir) and args.local_rank in [-1, 0]:
        os.makedirs(eval_output_dir)

    args.eval_batch_size = args.per_gpu_eval_batch_size * max(1, args.n_gpu)
    # Note that DistributedSampler samples randomly
    eval_sampler = SequentialSampler(eval_dataset) if args.local_rank == -1 else DistributedSampler(eval_dataset)
    eval_dataloader = DataLoader(eval_dataset, sampler=eval_sampler, batch_size=args.eval_batch_size)

    # multi-gpu eval
    if args.n_gpu > 1:
        model = torch.nn.DataParallel(model)
        super_model = torch.nn.DataParallel(super_model)

    # Eval!
    logger.info("***** Running evaluation {} *****".format(prefix))
    logger.info("  Num examples = %d", len(eval_dataset))
    logger.info("  Batch size = %d", args.eval_batch_size)
    nb_eval_steps = 0
    logits = []
    probs = []
    out_label_ids = []
    all_macs = []
    total_goto_super = 0
    start_time = time.time()
    for batch in tqdm(eval_dataloader, desc="Evaluating"):
        model.eval()
        super_model.eval()
        batch = tuple(t.to(args.device) for t in batch)

        with torch.no_grad():
            # go through the base model (lstm) first
            inputs = {
                'input_ids': batch[0],
                'attention_mask': batch[1],
                'labels': batch[3],
            }
            if args.model_type in ['bert', 'xlnet']:
                inputs['token_type_ids'] = batch[2]  # XLM, DistilBERT and RoBERTa don't use segment_ids
            if profile:
                macs, batch_logits = profile_macs(model, batch[:3])
                batch_logits = batch_logits[0]
                all_macs.append(macs)
            else:
                _, batch_logits = model(**inputs)
                batch_logits = batch_logits[0]
            batch_probs = F.softmax(batch_logits, dim=1)

            # decide which samples to send to super_model
            if args.random_multi_stage_proportion is None:
                goto_super = batch_probs.max(dim=1)[0] < args.conf_th  # whether to sent a sample to super_model
            else:
                # just the first few samples in the batch
                N = len(batch[0])
                goto_super = torch.tensor(
                    [(i/N < args.random_multi_stage_proportion) for i in range(N)]
                ).to(batch[0].device)
            goto_super_count = sum(goto_super)
            total_goto_super += goto_super_count

            # go through the super model (bert)
            if goto_super_count > 0:
                picked_inputs = {
                    'input_ids': batch[0][goto_super],
                    'attention_mask': batch[1][goto_super],
                    'labels': batch[3][goto_super],
                }
                if args.model_type in ['bert', 'xlnet']:
                    picked_inputs['token_type_ids'] = batch[2][goto_super]  # XLM, DistilBERT and RoBERTa don't use segment_ids
                if profile:
                    macs, super_batch_logits = profile_macs(
                        super_model,
                        (batch[0][goto_super], batch[1][goto_super], batch[2][goto_super])
                    )
                    super_batch_logits = super_batch_logits[0]
                    all_macs.append(macs)
                else:
                    _, super_batch_logits = super_model(**picked_inputs)
                super_batch_probs = F.softmax(super_batch_logits, dim=1)

        nb_eval_steps += 1

        logits.append(batch_logits[~goto_super].detach().cpu().numpy())
        probs.append(batch_probs[~goto_super].detach().cpu().numpy())
        out_label_ids.append(batch[3][~goto_super].detach().cpu().numpy())
        if goto_super_count > 0:
            logits.append(super_batch_logits.detach().cpu().numpy())
            probs.append(super_batch_probs.detach().cpu().numpy())
            out_label_ids.append(batch[3][goto_super].detach().cpu().numpy())

    eval_time = time.time() - start_time
    logits = np.concatenate(logits, axis=0)
    probs = np.concatenate(probs, axis=0)
    out_label_ids = np.concatenate(out_label_ids, axis=0)

    if args.output_mode == "selective-classification":
        th_results = {}  # threshold -> score
        for threshold in np.arange(0.7, 1.0, 0.001):
            preds = np.argmax(logits, axis=1)
            for i in range(len(preds)):
                if max(probs[i]) < threshold:
                    preds[i] = model.num_labels - 1
            th_results[threshold] = compute_metrics(eval_task, preds, out_label_ids)['acc']

        best_th, best_acc = 0, 0
        for th, acc in th_results.items():
            # print(th, acc)
            if acc > best_acc:
                best_th, best_acc = th, acc
        print({
            'best_th': best_th,
            'best_acc': best_acc,
        })

    if args.output_mode in ["classification", "selective-classification"]:
        preds = np.argmax(logits, axis=1)
    elif args.output_mode == "regression":
        preds = np.squeeze(logits)
    result = compute_metrics(eval_task, preds, out_label_ids)
    print_result = get_wanted_result(result)

    conf = probs.max(axis=1)
    risk_binary = (preds != out_label_ids).astype(int)
    rcc_auc_value = rcc_auc(conf, risk_binary)
    rpp_value = rpp(conf, risk_binary)

    result_dict = {
        "goto_super": float(total_goto_super) / len(eval_dataset),
        "eval_time": eval_time,
        "result": print_result,
        "rcc_auc": rcc_auc_value,
        "rpp": rpp_value,
        "macs(G)": sum(all_macs) / len(out_label_ids) / 1e9,
    }

    print(result_dict)
    results.update(result)

    save_fname = build_file(
        args.plot_data_dir,
        'saved_data',
        os.path.basename(args.data_dir),
        ('' if args.random_multi_stage_proportion is None else 'random-') + \
        args.multi_stage_base + '-' + args.train_routine + '.npy'
    )
    record_into_file(save_fname, args.conf_th, result_dict)

    output_eval_file = os.path.join(eval_output_dir, prefix, "eval_results.txt")
    with open(output_eval_file, "w") as writer:
        logger.info("***** Eval results {} *****".format(prefix))
        for key in sorted(result.keys()):
            logger.info("  %s = %s", key, str(result[key]))
            writer.write("%s = %s\n" % (key, str(result[key])))

    return results

def build_file(*args):
    save_fname = os.path.join(*args)
    if not os.path.exists(save_fname):
        os.makedirs(os.path.dirname(save_fname), exist_ok=True)
    return save_fname

def record_into_file(fname, key, value):
    if os.path.exists(fname):
        saved_dict = np.load(fname, allow_pickle=True).item()
    else:
        saved_dict = {}

    if key in saved_dict:
        print(key, 'already in the file!')
    else:
        saved_dict[key] = value
        np.save(fname, saved_dict)


#!/usr/bin/env python

import rospy
from std_msgs.msg import Empty
from sensor_msgs.msg import Image
from geometry_msgs.msg import Twist

import numpy as np
import math

import cv2
from cv_bridge import CvBridge, CvBridgeError
bridge = CvBridge()

global_image = None
global_rel_gt = None
save_images = False

# Constants
D_H_SHORT = 3.0
D_H_LONG = 9.0
D_ARROW = 30.0
D_RADIUS = 40.0

# Image size
IMG_WIDTH = 640
IMG_HEIGHT = 360

#################
# ROS functions #
#################
def image_callback(data):
    global global_image

    try:
        global_image = bridge.imgmsg_to_cv2(data, 'bgr8') # {'bgr8' or 'rgb8}
    except CvBridgeError as e:
        rospy.loginfo(e)

def rel_gt_callback(data):
    global global_rel_gt
    global_rel_gt = data

##################
# Help functions #
##################
def make_gaussian_blurry(image, blur_size):
    return cv2.GaussianBlur(image, (blur_size, blur_size), 0)

def make_circle_average_blurry(image, blur_size):
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (blur_size,blur_size))

    n_elements = np.float64(np.count_nonzero(kernel))
    kernel_norm = (kernel/n_elements)

    img_blurred = cv2.filter2D(image,-1,kernel_norm)
    return img_blurred

def hsv_save_image(image, label='image', is_gray=False):
    folder = './image_processing/'
    if save_images:
        if is_gray:
            cv2.imwrite(folder+label+".png", image)
        else:
            cv2.imwrite(folder+label+".png", cv2.cvtColor(image, cv2.COLOR_HSV2BGR))
    return image

def rgb_color_to_hsv(red, green, blue):
    bgr_color = np.uint8([[[blue,green,red]]])
    hsv_color = cv2.cvtColor(bgr_color,cv2.COLOR_BGR2HSV)
    return hsv_color[0][0].tolist()

def normalize_vector(vector):
    return vector / np.linalg.norm(vector)

def hsv_to_opencv_hsv(hue, saturation, value):
    """
    Function that takes in hue, saturation and value in the ranges
        hue: [0, 360] degrees,  saturation: [0, 100] %,     value: [0, 100] %
    and converts it to OpenCV hsv which operates with the ranges
        hue: [0, 180],          saturation: [0, 255],       value: [0, 255]
    """
    converting_constant = np.array([0.5, 2.55, 2.55])
    return np.array([ hue, saturation, value])*converting_constant

def draw_dot(img, position, color, size=3):
    if save_images:
        cX = np.int0(position[1])
        cY = np.int0(position[0])
        cv2.circle(img, (cX, cY), size, color, -1)

def draw_arrow(img, start, end):
    if save_images:
        return cv2.arrowedLine(img,
            (np.int0(start[1]), np.int0(start[0])),
            (np.int0(end[1]), np.int0(end[0])),
            color = (75,0,0), thickness = 1, tipLength = 0.4)

def calc_angle_between_vectors(vector_1, vector_2):
    v1_x = vector_1[0]
    v1_y = vector_1[1]

    v2_x = vector_2[0]
    v2_y = vector_2[1]

    return np.arctan2( v1_x*v2_y - v1_y*v2_x, v1_x*v2_x + v1_y*v2_y)

def get_normal_vector(bw_white_mask, corner_a, corner_b, is_short_side):
    hsv_normals = bw_white_mask.copy()

    vector_between_a_b = corner_a - corner_b

    vector_length = np.linalg.norm(vector_between_a_b)
    unit_vector_between_a_b = normalize_vector(vector_between_a_b)

    v_x, v_y = vector_between_a_b
    normal_unit_vector_left = normalize_vector(np.array([-v_y, v_x]))
    normal_unit_vector_right = normalize_vector(np.array([v_y, -v_x]))

    if is_short_side:
        check_length = vector_length / 3.0
        sign = 1 # Go outwards from the corners
    else:
        short_length = vector_length * D_H_SHORT/D_H_LONG

        check_length = short_length / 3.0
        sign = -1 # Go inwards from the corners

    check_left_a = np.int0(corner_a + \
        sign*unit_vector_between_a_b*check_length + \
        normal_unit_vector_left*check_length)
    check_left_b = np.int0(corner_b - \
        sign*unit_vector_between_a_b*check_length + \
        normal_unit_vector_left*check_length)

    check_right_a = np.int0(corner_a + \
        sign*unit_vector_between_a_b*check_length + \
        normal_unit_vector_right*check_length)
    check_right_b = np.int0(corner_b - \
        sign*unit_vector_between_a_b*check_length + \
        normal_unit_vector_right*check_length)

    value_left_a = bw_white_mask[check_left_a[0]][check_left_a[1]]
    value_left_b = bw_white_mask[check_left_b[0]][check_left_b[1]]
    value_right_a = bw_white_mask[check_right_a[0]][check_right_a[1]]
    value_right_b = bw_white_mask[check_right_b[0]][check_right_b[1]]

    avr_left = value_left_a/2.0 + value_left_b/2.0
    avr_right = value_right_a/2.0 + value_right_b/2.0

    draw_dot(hsv_normals, check_left_a, (225))
    draw_dot(hsv_normals, check_left_b, (225))
    draw_dot(hsv_normals, check_right_a, (75))
    draw_dot(hsv_normals, check_right_b, (75))

    # hsv_save_image(hsv_normals, "5_normals", is_gray=True)

    if avr_left > avr_right:
        return normal_unit_vector_left
    else:
        return normal_unit_vector_right

# Colors to draw with
HSV_RED_COLOR = rgb_color_to_hsv(255,0,0)
HSV_BLUE_COLOR = rgb_color_to_hsv(0,0,255)
HSV_BLACK_COLOR = rgb_color_to_hsv(0,0,0)
HSV_YELLOW_COLOR = [30, 255, 255]
HSV_LIGHT_ORANGE_COLOR = [15, 255, 255]

HSV_REAL_WHITE = [0, 0, 74]
HSV_REAL_ORANGE = [36, 100, 74]
HSV_REAL_GREEN = [120, 100, 30]

# Setup for color finding
HUE_MARGIN = 15
SAT_MARGIN = 15
VAL_MARGIN = 15

# White
HUE_LOW_WHITE = 0
HUE_HIGH_WHITE = 360
SAT_LOW_WHITE = max(0, HSV_REAL_WHITE[1] - SAT_MARGIN)
SAT_HIGH_WHITE = min(100, HSV_REAL_WHITE[1] + SAT_MARGIN)
VAL_LOW_WHITE = max(0, HSV_REAL_WHITE[2] - VAL_MARGIN)
VAL_HIGH_WHITE = min(100, HSV_REAL_WHITE[2] + VAL_MARGIN)

# Orange
HUE_LOW_ORANGE = max(0, HSV_REAL_ORANGE[0] - HUE_MARGIN)
HUE_HIGH_ORANGE = min(360, HSV_REAL_ORANGE[0] + HUE_MARGIN)
SAT_LOW_ORANGE = max(0, HSV_REAL_ORANGE[1] - SAT_MARGIN)
SAT_HIGH_ORANGE = min(100, HSV_REAL_ORANGE[1] + SAT_MARGIN)
VAL_LOW_ORANGE = max(0, HSV_REAL_ORANGE[2] - VAL_MARGIN)
VAL_HIGH_ORANGE = min(100, HSV_REAL_ORANGE[2] + VAL_MARGIN)

# Green
HUE_LOW_GREEN = max(0, HSV_REAL_GREEN[0] - HUE_MARGIN)
HUE_HIGH_GREEN = min(360, HSV_REAL_GREEN[0] + HUE_MARGIN)
SAT_LOW_GREEN = max(0, HSV_REAL_GREEN[1] - SAT_MARGIN)
SAT_HIGH_GREEN = min(100, HSV_REAL_GREEN[1] + SAT_MARGIN)
VAL_LOW_GREEN = max(0, HSV_REAL_GREEN[2] - VAL_MARGIN)
VAL_HIGH_GREEN = min(100, HSV_REAL_GREEN[2] + VAL_MARGIN)

##################
# Main functions #
##################
def get_mask(hsv, hue_low, hue_high, sat_low, sat_high, val_low, val_high):
    lower_color = hsv_to_opencv_hsv(hue_low, sat_low, val_low)
    upper_color = hsv_to_opencv_hsv(hue_high, sat_high, val_high)

    mask = cv2.inRange(hsv, lower_color, upper_color)

    mask_x, mask_y = np.where(mask==255)
    if len(mask_x) == 0: # No color visible
        return None
    else:
        return mask

def get_white_mask(hsv):
    return get_mask(hsv, HUE_LOW_WHITE, HUE_HIGH_WHITE,SAT_LOW_WHITE, SAT_HIGH_WHITE, VAL_LOW_WHITE, VAL_HIGH_WHITE)

def get_orange_mask(hsv):
    return get_mask(hsv, HUE_LOW_ORANGE, HUE_HIGH_ORANGE, SAT_LOW_ORANGE, SAT_HIGH_ORANGE, VAL_LOW_ORANGE, VAL_HIGH_ORANGE)

def get_green_mask(hsv):
    return get_mask(hsv, HUE_LOW_GREEN, HUE_HIGH_GREEN, SAT_LOW_GREEN, SAT_HIGH_GREEN, VAL_LOW_GREEN, VAL_HIGH_GREEN)

def flood_fill(img, start=(0,0)):
    h,w = img.shape
    seed = start

    mask = np.zeros((h+2,w+2),np.uint8) # Adding a padding of 1

    if img[start[1]][start[0]] == 255: # If the starting point is already filled, return empty mask
        mask = mask[1:h+1,1:w+1] # Removing the padding
        return mask

    floodflags = 8
    floodflags |= cv2.FLOODFILL_MASK_ONLY
    floodflags |= (255 << 8)

    num,img,mask,rect = cv2.floodFill(img, mask, seed, (255,0,0), (10,)*3, (10,)*3, floodflags)
    mask = mask[1:h+1,1:w+1] # Removing the padding

    return mask

def get_pixels_inside_green(hsv):
    """
        Function that finds the green in an image, make a bounding box around it,
        fits an ellipse in the bounding box
        and paints everything outside the ellipse in black.

        Returns the painted image and a boolean stating wheather any green was found.
     """

    bw_green_mask = get_green_mask(hsv)
    if bw_green_mask is None:
        return hsv

    green_x, green_y = np.where(bw_green_mask==255)
    x_min = np.amin(green_x)
    x_max = np.amax(green_x)
    y_min = np.amin(green_y)
    y_max = np.amax(green_y)

    center_x = np.int0((x_min + x_max) / 2.0)
    center_y = np.int0((y_min + y_max) / 2.0)

    l_x = np.int0(x_max - center_x)+1
    l_y = np.int0(y_max - center_y)+1

    bw_ellipse_mask = np.zeros((IMG_HEIGHT, IMG_WIDTH))
    cv2.ellipse(img=bw_ellipse_mask, center=(center_y, center_x),
        axes=(l_y, l_x), angle=0, startAngle=0, endAngle=360,
        color=(255), thickness=-1, lineType=8, shift=0)

    hsv_ellipse = hsv.copy()
    hsv_ellipse[bw_ellipse_mask==0] = HSV_BLACK_COLOR

    return hsv_ellipse

def get_pixels_inside_orange(hsv):
    """
        Function that finds the orange in an image, make a bounding box around it,
        and paints everything outside the box in black.

        Returns the painted image and a boolean stating wheather any orange was found.
     """
    bw_orange_mask = get_orange_mask(hsv)
    if bw_orange_mask is None:
        return hsv
    hsv_save_image(bw_orange_mask, "2b_orange_mask", is_gray=True)

    orange_x, orange_y = np.where(bw_orange_mask==255)
    x_min = np.amin(orange_x)
    x_max = np.amax(orange_x)
    y_min = np.amin(orange_y)
    y_max = np.amax(orange_y)

    hsv_inside_orange = hsv.copy()
    hsv_inside_orange[0:x_min,] = HSV_BLACK_COLOR
    hsv_inside_orange[x_max+1:,] = HSV_BLACK_COLOR
    hsv_inside_orange[:,0:y_min] = HSV_BLACK_COLOR
    hsv_inside_orange[:,y_max+1:] = HSV_BLACK_COLOR
    hsv_save_image(hsv_inside_orange, '3_inside_orange')

    return hsv_inside_orange

def find_white_centroid(hsv):
    hsv_inside_orange = get_pixels_inside_orange(hsv)

    bw_white_mask = get_white_mask(hsv_inside_orange)
    if bw_white_mask is None:
        rospy.loginfo("bw_white_mask is none")
        return None
    bw_white_mask = make_gaussian_blurry(bw_white_mask, 5)

    M = cv2.moments(bw_white_mask) # calculate moments of binary image

    # calculate x,y coordinate of center
    cX = int(M["m10"] / M["m00"])
    cY = int(M["m01"] / M["m00"])

    # Returnes the transposed point,
    #  because of difference from OpenCV axis
    return np.array([cY, cX])

def find_harris_corners(img, block_size):
    """ Using sub-pixel method from OpenCV """
    dst = cv2.cornerHarris(img, block_size, 3, 0.04)
    dst = cv2.dilate(dst, None)
    ret, dst = cv2.threshold(dst,0.01*dst.max(), 255, 0)
    dst = np.uint8(dst)

    # find centroids
    ret, labels, stats, centroids = cv2.connectedComponentsWithStats(dst)

    # define the criteria to stop and refine the corners
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 100, 0.001)
    corners = cv2.cornerSubPix(img,np.float32(centroids),(5,5),(-1,-1),criteria)

    # Flip axis
    corners[:,[0, 1]] = corners[:,[1, 0]]

    return corners

def clip_corners_on_border(corners, border_size):
    # Filter the corners
    number_of_corners = len(corners)
    x_min = np.array([border_size]*number_of_corners)
    x_max = np.array([IMG_HEIGHT - border_size]*number_of_corners)
    y_min = np.array([border_size]*number_of_corners)
    y_max = np.array([IMG_WIDTH - border_size]*number_of_corners)

    corner_x = np.int0(corners[:,0])
    corner_y = np.int0(corners[:,1])

    # Keep corners within the border limit
    corners_clipped_on_border = corners[
        np.logical_and(
            np.logical_and(
                np.greater(corner_x, x_min),            # Add top limit
                np.less(corner_x, x_max)),              # Add bottom limit
            np.logical_and(
                np.greater(corner_y, y_min),            # Add left limit
                np.less(corner_y, y_max))               # Add right limit
        )
    ]
    corner_x = np.int0(corners_clipped_on_border[:,0])
    corner_y = np.int0(corners_clipped_on_border[:,1])

    if np.ndim(corner_x) == 0:
        number_of_corners = 1
        corners = np.array([[corner_x, corner_y]])
    else:
        corners = np.stack((corner_x, corner_y), axis=1)
        number_of_corners = len(corners)
    if number_of_corners == 0:
        return None
    else:
        return corners

def clip_corners_on_intensity(corners, img, average_filter_size):
    """
        Filter out the corners that belong to a right-angled corner
        i.e. corners with a mean intensity value around 255/4~64
    """
    value_per_degree = 255.0/360.0
    min_degree, max_degree = 60, 120 # +- 30 from 90 degrees

    # Since 255 is white and 0 is black, subtract from 255
    # to get black intensity instead of white intensity
    min_average_intensity = 255 - max_degree*value_per_degree
    max_average_intensity = 255 - min_degree*value_per_degree

    number_of_corners = len(corners)
    min_intensity = np.array([min_average_intensity]*number_of_corners)
    max_intensity = np.array([max_average_intensity]*number_of_corners)

    img_average_intensity = make_circle_average_blurry(img, average_filter_size)

    corner_x = np.int0(corners[:,0])
    corner_y = np.int0(corners[:,1])

    corners_clipped_on_intensity = corners[
        np.logical_and(
            np.greater(                                                 # Add top limit
                img_average_intensity[corner_x,corner_y],
                min_intensity),
            np.less(                                                    # Add bottom limit
                img_average_intensity[corner_x,corner_y],
                max_intensity)
        )
    ]
    corner_x = np.int0(corners_clipped_on_intensity[:,0])
    corner_y = np.int0(corners_clipped_on_intensity[:,1])

    if np.ndim(corner_x) == 0:
        corners = np.array([[corner_x, corner_y]])
        intensities = np.array([img_average_intensity[corner_x, corner_y]])
        number_of_corners = 1
    else:
        corners = np.stack((corner_x, corner_y), axis=1)
        intensities = np.array(img_average_intensity[corner_x, corner_y])
        number_of_corners = len(corners)

    if number_of_corners == 0:
        return None, None
    else:
        return corners, intensities

def find_right_angled_corners(img):
    # Parameters ###########################
    average_filter_size = 19 # 19
    ignore_border_size = 3
    corner_harris_block_size = 4

    # Define valid intensity range for the median of a corner
    min_intensity_average = 170
    max_intensity_average = 240
    ########################################

    corners = find_harris_corners(img, corner_harris_block_size)
    corners = clip_corners_on_border(corners, ignore_border_size)
    if corners is None:
        return None, None

    corners, intensities = clip_corners_on_intensity(corners, img, average_filter_size)
    if corners is None:
        return None, None

    return corners, intensities

def find_orange_arrowhead(hsv):
    bw_orange_mask = get_orange_mask(hsv)
    if bw_orange_mask is None:
        return None

    bw_orange_mask = make_gaussian_blurry(bw_orange_mask, 5)
    bw_orange_mask_inverted = cv2.bitwise_not(bw_orange_mask)

    hsv_save_image(bw_orange_mask_inverted, "0_orange_mask_inverted", is_gray=True)

    orange_corners, intensities = find_right_angled_corners(bw_orange_mask_inverted)
    if orange_corners is None:
        return None

    number_of_corners_found = len(orange_corners)

    value_per_degree = 255.0/360.0
    ideal_angle = 90
    ideal_intensity = 255-ideal_angle*value_per_degree
    ideal_intensities = np.array([ideal_intensity]*number_of_corners_found)

    diff_intensities = np.absolute(np.array(ideal_intensities-intensities))

    if number_of_corners_found == 1:
        return orange_corners[0]
    elif number_of_corners_found > 1:
        # Too many orange corners found, choose the best
        best_corner_id = np.argmin(diff_intensities)
        best_corner = orange_corners[best_corner_id]
        return best_corner
    else:
        # No corners found
        return None

def calculate_position(center_px, radius_px):
    focal_length = 374.67
    real_radius = 375 # mm (750mm in diameter / 2)

    # Center of image
    x_0 = IMG_HEIGHT/2.0
    y_0 = IMG_WIDTH/2.0

    # Find distances from center of image to center of LP
    d_x = x_0 - center_px[0]
    d_y = y_0 - center_px[1]

    est_z = real_radius*focal_length / radius_px # - 59.4 # (adjustment)

    # Camera is placed 150 mm along x-axis of the drone
    # Since the camera is pointing down, the x and y axis of the drone
    # is the inverse of the x and y axis of the camera
    est_x = -(est_z * d_x / focal_length) - 150
    est_y = -(est_z * d_y / focal_length)

    return np.array([est_x, est_y, est_z])

def fit_ellipse(points):
    x = points[1]
    y = IMG_HEIGHT-points[0]

    D11 = np.square(x)
    D12 = x*y
    D13 = np.square(y)
    D1 = np.array([D11, D12, D13]).T
    D2 = np.array([x, y, np.ones(x.shape[0])]).T

    S1 = np.dot(D1.T,D1)
    S2 = np.dot(D1.T,D2)
    S3 = np.dot(D2.T,D2)

    try:
        inv_S3 = np.linalg.inv(S3)
    except np.linalg.LinAlgError:
        print("fit_ellipse(): Got singular matrix")
        return None

    T = - np.dot(inv_S3, S2.T) # for getting a2 from a1

    M = S1 + np.dot(S2, T)

    C1 = np.array([
        [0, 0, 0.5],
        [0, -1, 0],
        [0.5, 0, 0]
    ])

    M = np.dot(C1, M) # This premultiplication can possibly be made more efficient

    eigenvalues, eigenvectors = np.linalg.eig(M)
    cond = 4*eigenvectors[0]*eigenvectors[2] - np.square(eigenvectors[0])
    a1 = eigenvectors[:,cond > 0]

    # Choose the first if there are two eigenvectors with cond > 0
    # NB! I am not sure if this is always correct
    if a1.shape[1] > 1:
        a1 = np.array([a1[:,0]]).T

    if a1.shape != (3,1): # Make sure a1 has content
        print("fit_ellipse(): a1 not OK")
        return None

    a = np.concatenate((a1, np.dot(T, a1)))[:,0] # Choose the inner column with [:,0]

    if np.any(np.iscomplex(a)):
        print("Found complex number")
        return None
    else:
        return a

def get_ellipse_parameters(green_ellipse):
    edges = cv2.Canny(green_ellipse,100,200)
    result = np.where(edges == 255)

    ellipse = fit_ellipse(result)

    if ellipse is None:
        return None

    A = ellipse[0]
    B = ellipse[1]
    C = ellipse[2]
    D = ellipse[3]
    E = ellipse[4]
    F = ellipse[5]

    if B**2 - 4*A*C >= 0:
        print("get_ellipse_parameters(): Shape found is not an ellipse")
        return None

    inner_square = math.sqrt( (A-C)**2 + B**2)
    outside = 1.0 / (B**2 - 4*A*C)
    a = outside * math.sqrt(2*(A*E**2 + C*D**2 - B*D*E + (B**2 - 4*A*C)*F) * ( (A+C) + inner_square))
    b = outside * math.sqrt(2*(A*E**2 + C*D**2 - B*D*E + (B**2 - 4*A*C)*F) * ( (A+C) - inner_square))

    x_raw = (2.0*C*D - B*E) / (B*B - 4.0*A*C)
    y_raw = (2.0*A*E - B*D) / (B*B - 4.0*A*C)

    x_0 = (IMG_HEIGHT - 1) - y_raw
    y_0 = x_raw

    ellipse_and_a_b = np.array([x_0,y_0,a,b])
    return ellipse_and_a_b

def evaluate_ellipse(hsv):
    """ Use the green ellipse to find:
        center, radius, angle
    """
    bw_green_mask = get_green_mask(hsv)
    if bw_green_mask is None:
        return None, None, None

    hsv_save_image(bw_green_mask, "2_green_mask", is_gray=True)

    top_border =    bw_green_mask[0,:]
    bottom_border = bw_green_mask[IMG_HEIGHT-1,:]
    left_border =   bw_green_mask[:,0]
    right_border =  bw_green_mask[:,IMG_WIDTH-1]

    sum_top_border = np.sum(top_border) + \
        np.sum(bottom_border) + \
        np.sum(left_border) + \
        np.sum(right_border)

    if sum_top_border != 0:
        # Then the green ellipse is toughing the border
        return None, None, None

    bw_green_ellipse = flood_fill(bw_green_mask, start=(0,0))
    hsv_save_image(bw_green_ellipse, "3_green_ellipse", is_gray=True)

    ellipse_parameters = get_ellipse_parameters(bw_green_ellipse)
    if ellipse_parameters is None:
        return None, None, None

    # Choose the largest of the a and b parameter for the radius
    # Choose angle = 0 since it is not possible to estimate from the ellipse
    center_px = ellipse_parameters[0:2]
    radius_px = np.amax(np.abs(ellipse_parameters[2:4]))
    angle = 0

    hsv_canvas_ellipse = hsv.copy()
    draw_dot(hsv_canvas_ellipse, center_px, HSV_BLUE_COLOR)
    hsv_save_image(hsv_canvas_ellipse, "4_canvas_ellipse") #, is_gray=True)

    return center_px, radius_px, angle

def evaluate_arrow(hsv):
    """ Use the arrow to find:
        center, radius, angle
    """
    center_px = find_white_centroid(hsv)
    arrowhead_px = find_orange_arrowhead(hsv)

    if (center_px is not None) and (arrowhead_px is not None):

        arrow_vector = np.array(arrowhead_px - center_px)
        arrow_unit_vector = normalize_vector(arrow_vector)
        ref_vector = np.array([0,1])

        angle = calc_angle_between_vectors(arrow_vector, ref_vector)

        arrow_length_px = np.linalg.norm(arrow_vector)
        # Use known relation between the real radius and the real arrow length
        # to find the radius length in pixels

        radius_length_px = arrow_length_px * D_RADIUS / D_ARROW

        hsv_canvas_arrow = hsv.copy()
        draw_dot(hsv_canvas_arrow, center_px, HSV_RED_COLOR)
        draw_dot(hsv_canvas_arrow, arrowhead_px, HSV_RED_COLOR)
        hsv_save_image(hsv_canvas_arrow, "4_canvas_arrow")

        return center_px, radius_length_px, angle

    else:
        return None, None, None

def get_relevant_corners(inner_corners):
    n_inner_corners = len(inner_corners)

    # For the first corner, chose the corner closest to the top
    # This will belong to the top cross-bar
    top_corner_id = np.argmin(inner_corners[:,0]) # Find lowest x-index
    top_corner = inner_corners[top_corner_id]
    top_corner_stack = np.array([top_corner]*(n_inner_corners-1))
    rest_corners = np.delete(inner_corners, top_corner_id, 0)
    dist = np.linalg.norm(rest_corners - top_corner_stack, axis=1)

    # For the second corner, chose the corner closest to top corner
    top_corner_closest_id = np.argmin(dist)
    top_corner_closest = rest_corners[top_corner_closest_id]

    relevant_corners = np.stack((top_corner, top_corner_closest), axis=0)

    return relevant_corners

def evaluate_inner_corners(hsv):
    """ Use the inner corners to find:
        center, radius, angle
    """
    hsv_canvas = hsv.copy()

    bw_white_mask = get_white_mask(hsv)
    if bw_white_mask is None:
        return None, None, None
    bw_white_mask = make_gaussian_blurry(bw_white_mask, 5)

    hsv_save_image(bw_white_mask, "0_white_only", is_gray=True)

    inner_corners, intensities = find_right_angled_corners(bw_white_mask)

    average_filter_size = 19
    img_average_intensity = make_circle_average_blurry(bw_white_mask, average_filter_size)

    if (inner_corners is not None):
        n_inner_corners = len(inner_corners)
        if (n_inner_corners > 1) and (n_inner_corners <= 5):
            unique_corners = np.vstack({tuple(row) for row in inner_corners}) # Removes duplicate corners

            for corner in unique_corners:
                draw_dot(hsv_canvas, corner, HSV_YELLOW_COLOR)

            corner_a, corner_b = get_relevant_corners(unique_corners)
            c_m = (corner_a + corner_b)/2.0 # Finds the mid point between the corners
            c_m_value = img_average_intensity[np.int0(c_m[0])][np.int0(c_m[1])]

            draw_dot(hsv_canvas, corner_a, HSV_RED_COLOR)
            draw_dot(hsv_canvas, corner_b, HSV_LIGHT_ORANGE_COLOR)
            draw_dot(hsv_canvas, c_m, HSV_BLUE_COLOR)
            hsv_save_image(hsv_canvas, "3_canvas")

            if c_m_value > 190: # The points are on a short side
                is_short_side = True
                normal_vector = get_normal_vector(bw_white_mask, corner_a, corner_b, is_short_side)
                normal_unit_vector = normalize_vector(normal_vector)

                length_short_side = np.linalg.norm(corner_a - corner_b)
                length_long_side = length_short_side * D_H_LONG/D_H_SHORT
                length_to_center = - length_long_side / 2.0
                length_radius = length_short_side * D_RADIUS/D_H_SHORT

                forward_unit_vector = normal_unit_vector

            else: # The points are on a long side
                is_short_side = False
                normal_vector = get_normal_vector(bw_white_mask, corner_a, corner_b, is_short_side)
                normal_unit_vector = normalize_vector(normal_vector)

                length_long_side = np.linalg.norm(corner_a - corner_b)
                length_short_side = length_long_side * D_H_SHORT/D_H_LONG
                length_to_center = length_short_side / 2.0
                length_radius = length_long_side * D_RADIUS/D_H_LONG

                forward_unit_vector = normalize_vector(corner_a - corner_b)

            end = c_m + forward_unit_vector*10
            draw_arrow(hsv_canvas, c_m, end)

            center = c_m + normal_unit_vector*length_to_center
            draw_dot(hsv_canvas, center, HSV_BLUE_COLOR)

            # hsv_save_image(hsv_canvas, "3_canvas")

            neg_x_axis = np.array([-1,0])
            angle = calc_angle_between_vectors(forward_unit_vector, neg_x_axis)

            hsv_canvas_inner_corners = hsv.copy()
            draw_dot(hsv_canvas_inner_corners, center, HSV_LIGHT_ORANGE_COLOR)
            hsv_save_image(hsv_canvas_inner_corners, "4_canvas_inner_corners")

            return center, length_radius, angle

    return None, None, None

def ros_run(hsv, count):
    hsv_save_image(hsv, '0_hsv')

    white_mask = get_white_mask(hsv)
    if white_mask is not None:
        hsv_save_image(white_mask, '1_white_mask', is_gray=True)

    orange_mask = get_orange_mask(hsv)
    if orange_mask is not None:
        hsv_save_image(orange_mask, '1_orange_mask', is_gray=True)

    green_mask = get_green_mask(hsv)
    if green_mask is not None:
        hsv_save_image(green_mask, '1_green_mask', is_gray=True)

    hsv_inside_green = get_pixels_inside_green(hsv)

    center_px_from_ellipse, radius_length_px_from_ellipse, angle_from_ellipse = evaluate_ellipse(hsv)
    center_px_from_arrow, radius_length_px_from_arrow, angle_from_arrow = evaluate_arrow(hsv) # or use hsv_inside_green
    center_px_from_inner_corners, radius_px_length_from_inner_corners, angle_from_inner_corners = evaluate_inner_corners(hsv_inside_green)

    hsv_canvas_all = hsv.copy()

    ############
    # Method 1 #
    ############
    if (center_px_from_ellipse is not None):
        center_px, radius_length_px, angle_rad = center_px_from_ellipse, radius_length_px_from_ellipse, angle_from_ellipse
        est_ellipse_x, est_ellipse_y, est_ellipse_z = calculate_position(center_px, radius_length_px)
        est_ellipse_angle = np.degrees(angle_rad)

        draw_dot(hsv_canvas_all, center_px, HSV_BLUE_COLOR, size=5)
    else:
        est_ellipse_x, est_ellipse_y, est_ellipse_z, est_ellipse_angle = 0.0, 0.0, 0.0, 0.0

    ############
    # Method 2 #
    ############
    if (center_px_from_arrow is not None):
        center_px, radius_length_px, angle_rad = center_px_from_arrow, radius_length_px_from_arrow, angle_from_arrow
        est_arrow_x, est_arrow_y, est_arrow_z = calculate_position(center_px, radius_length_px)
        est_arrow_angle = np.degrees(angle_rad)

        draw_dot(hsv_canvas_all, center_px, HSV_RED_COLOR, size=4)
    else:
        est_arrow_x, est_arrow_y, est_arrow_z, est_arrow_angle = 0.0, 0.0, 0.0, 0.0

    ############
    # Method 3 #
    ############
    if (center_px_from_inner_corners is not None):
        center_px, radius_length_px, angle_rad = center_px_from_inner_corners, radius_px_length_from_inner_corners, angle_from_inner_corners
        est_inner_corners_x, est_inner_corners_y, est_inner_corners_z = calculate_position(center_px, radius_length_px)
        est_inner_corners_angle = np.degrees(angle_rad)

        draw_dot(hsv_canvas_all, center_px, HSV_LIGHT_ORANGE_COLOR, size=3)
    else:
        est_inner_corners_x, est_inner_corners_y, est_inner_corners_z, est_inner_corners_angle = 0.0, 0.0, 0.0, 0.0

    hsv_save_image(hsv_canvas_all, "5_canvas_all_"+str(count))

    result = np.array([
        [est_ellipse_x, est_ellipse_y, est_ellipse_z, est_ellipse_angle],
        [est_arrow_x, est_arrow_y, est_arrow_z, est_arrow_angle],
        [est_inner_corners_x, est_inner_corners_y, est_inner_corners_z, est_inner_corners_angle]
    ])

    return result

def rel_gt_converter(rel_gt):
    """ Convert from twist message data to x, y, z, yaw in mm and degrees """
    gt_x = rel_gt.linear.x * 1000
    gt_y = rel_gt.linear.y * 1000
    gt_z = rel_gt.linear.z * 1000
    # yaw = -np.degrees(rel_gt.angular.z) - 90

    # if yaw < -180:
    #     gt_yaw = 360 + yaw
    # else:
    #     gt_yaw = yaw

    # gt_x = rel_gt.linear.x
    # gt_y = rel_gt.linear.y
    # gt_z = rel_gt.linear.z
    gt_yaw = rel_gt.angular.z

    return np.array([[gt_x, gt_y, gt_z, gt_yaw]])

def filter_estimate(estimate, estimate_history, median_filter_size, average_filter_size):
    """
        Filters the estimate with a sliding window median and average filter.
    """

    estimate_history = np.concatenate((estimate_history[1:], [estimate]))

    strides = np.array(
        [estimate_history[i:median_filter_size+i] for i in range(average_filter_size)]
    )

    median_filtered = np.median(strides, axis = 1)
    average_filtered = np.average(median_filtered[-average_filter_size:], axis=0)

    return average_filtered, estimate_history

def main():
    rospy.init_node('cv_module', anonymous=True)

    rospy.Subscriber('/ardrone/bottom/image_raw', Image, image_callback)
    rospy.Subscriber('/drone_pose', Twist, rel_gt_callback)

    pub_heartbeat = rospy.Publisher("/heartbeat", Empty, queue_size=10)

    pub_est_ellipse = rospy.Publisher("/estimate/ellipse", Twist, queue_size=10)
    pub_est_ellipse_filtered = rospy.Publisher("/estimate_filtered/ellipse", Twist, queue_size=10)

    pub_est_arrow = rospy.Publisher("/estimate/arrow", Twist, queue_size=10)
    pub_est_arrow_filtered = rospy.Publisher("/estimate_filtered/arrow", Twist, queue_size=10)

    pub_est_corners = rospy.Publisher("/estimate/corners", Twist, queue_size=10)
    pub_est_corners_filtered = rospy.Publisher("/estimate_filtered/corners", Twist, queue_size=10)

    est_ellipse_msg = Twist()
    est_filtered_ellipse_msg = Twist()

    est_arrow_msg = Twist()
    est_filtered_arrow_msg = Twist()

    est_corners_msg = Twist()
    est_filtered_corners_msg = Twist()

    heartbeat_msg = Empty()

    # Set up filter
    median_filter_size = 3
    average_filter_size = 3

    estimate_history_size = median_filter_size + average_filter_size - 1
    estimate_history_ellipse = np.zeros((estimate_history_size,4))
    estimate_history_arrow = np.zeros((estimate_history_size,4))
    estimate_history_corners = np.zeros((estimate_history_size,4))

    rospy.loginfo("Starting CV module")

    count = 0
    rate = rospy.Rate(20) # Hz
    while not rospy.is_shutdown():

        if (global_image is not None):
            pub_heartbeat.publish(heartbeat_msg)

            hsv = cv2.cvtColor(global_image, cv2.COLOR_BGR2HSV) # convert to HSV
            est = ros_run(hsv, count)
            # gt = rel_gt_converter(global_rel_gt)
            # result = np.concatenate((est, gt))

            est_ellipse = est[0]
            est_arrow = est[1]
            est_corners = est[2]

            est_ellipse_filtered, estimate_history_ellipse = filter_estimate(est_ellipse, estimate_history_ellipse, median_filter_size, average_filter_size)
            est_arrow_filtered, estimate_history_arrow = filter_estimate(est_arrow, estimate_history_arrow, median_filter_size, average_filter_size)
            est_corners_filtered, estimate_history_corners = filter_estimate(est_corners, estimate_history_corners, median_filter_size, average_filter_size)

            # Publish the results #
            # Ellipse
            est_ellipse_msg.linear.x = est_ellipse[0] / 1000.0
            est_ellipse_msg.linear.y = est_ellipse[1] / 1000.0
            est_ellipse_msg.linear.z = est_ellipse[2] / 1000.0
            est_ellipse_msg.angular.z = est_ellipse[3]
            pub_est_ellipse.publish(est_ellipse_msg)

            est_filtered_ellipse_msg.linear.x = est_ellipse_filtered[0] / 1000.0
            est_filtered_ellipse_msg.linear.y = est_ellipse_filtered[1] / 1000.0
            est_filtered_ellipse_msg.linear.z = est_ellipse_filtered[2] / 1000.0
            est_filtered_ellipse_msg.angular.z = est_ellipse_filtered[3]
            pub_est_ellipse_filtered.publish(est_filtered_ellipse_msg)

            # Arrow
            est_arrow_msg.linear.x = est_arrow[0] / 1000.0
            est_arrow_msg.linear.y = est_arrow[1] / 1000.0
            est_arrow_msg.linear.z = est_arrow[2] / 1000.0
            est_arrow_msg.angular.z = est_arrow[3]
            pub_est_arrow.publish(est_arrow_msg)

            est_filtered_arrow_msg.linear.x = est_arrow_filtered[0] / 1000.0
            est_filtered_arrow_msg.linear.y = est_arrow_filtered[1] / 1000.0
            est_filtered_arrow_msg.linear.z = est_arrow_filtered[2] / 1000.0
            est_filtered_arrow_msg.angular.z = est_arrow_filtered[3]
            pub_est_arrow_filtered.publish(est_filtered_arrow_msg)

            # Corners
            est_corners_msg.linear.x = est_corners[0] / 1000.0
            est_corners_msg.linear.y = est_corners[1] / 1000.0
            est_corners_msg.linear.z = est_corners[2] / 1000.0
            est_corners_msg.angular.z = est_corners[3]
            pub_est_corners.publish(est_corners_msg)

            est_filtered_corners_msg.linear.x = est_corners_filtered[0] / 1000.0
            est_filtered_corners_msg.linear.y = est_corners_filtered[1] / 1000.0
            est_filtered_corners_msg.linear.z = est_corners_filtered[2] / 1000.0
            est_filtered_corners_msg.angular.z = est_corners_filtered[3]
            pub_est_corners_filtered.publish(est_filtered_corners_msg)

            count += 1
        else:
            rospy.loginfo("Waiting for image")

        rate.sleep()

if __name__ == '__main__':
    main()


import numpy
def main():
    print("Operaciones Básicas:\n", "1. Suma\n", "2. Resta\n", "3. Multiplicacion\n", "4. Division\n", "5. Potencias\n", "6. Raiz Cuadrada\n")
    elejir = int(input("" ))
    if (elejir == 1):
        arr = []
        a = int(input("Ingrese cuantos valores va a sumar: "))
        for i in range(a):
            arr.append(float(input("Ingrese valor: ")))
        arr = numpy.array(arr)
        print(numpy.sum(arr, axis=0))

main()


import cups, pprint, os, tempfile, re, string
import locale
from . import _debugprint
from . import config
from functools import reduce

class Printer:
    _flags_blacklist = ["options", "local"]

    def __init__(self, name, connection, **kw):
        """
        @param name: printer name
        @type name: string
        @param connection: CUPS connection
        @type connection: CUPS.Connection object
        @param kw: printer attributes
        @type kw: dict indexed by string
        """
        self.name = name
        self.connection = connection
        self.class_members = []
        have_kw = len (kw) > 0
        fetch_attrs = True
        if have_kw:
            self.update (**kw)
            if self.is_class:
                fetch_attrs = True
            else:
                fetch_attrs = False

        if fetch_attrs:
            self.getAttributes ()

        self._ppd = None # load on demand

    def __del__ (self):
        if self._ppd is not None:
            os.unlink(self._ppd)

    def __repr__ (self):
        return "<cupshelpers.Printer \"%s\">" % self.name

    def _expand_flags(self):

        def _ascii_lower(str):
            return str.lower();

        prefix = "CUPS_PRINTER_"
        prefix_length = len(prefix)

        # loop over cups constants
        for name in cups.__dict__:
            if name.startswith(prefix):
                attr_name = \
                    _ascii_lower(name[prefix_length:])
                if attr_name in self._flags_blacklist: continue
                if attr_name == "class": attr_name = "is_class"
                # set as attribute
                setattr(self, attr_name,
                        bool(self.type & getattr(cups, name)))

    def update(self, **kw):
        """
        Update object from printer attributes.

        @param kw: printer attributes
        @type kw: dict indexed by string
        """
        self.state = kw.get('printer-state', 0)
        self.enabled = self.state != cups.IPP_PRINTER_STOPPED
        self.device_uri = kw.get('device-uri', "")
        self.info = kw.get('printer-info', "")
        self.is_shared = kw.get('printer-is-shared', None)
        self.location = kw.get('printer-location', "")
        self.make_and_model = kw.get('printer-make-and-model', "")
        self.type = kw.get('printer-type', 0)
        self.uri_supported = kw.get('printer-uri-supported', "")
        if type (self.uri_supported) != list:
            self.uri_supported = [self.uri_supported]
        self._expand_flags()
        if self.is_shared is None:
            self.is_shared = not self.not_shared
        del self.not_shared
        self.class_members = kw.get('member-names', [])
        if type (self.class_members) != list:
            self.class_members = [self.class_members]
        self.class_members.sort ()
        self.other_attributes = kw

    def getAttributes(self):
        """
        Fetch further attributes for the printer.

        Normally only a small set of attributes is fetched.  This
        method is for fetching more.
        """
        attrs = self.connection.getPrinterAttributes(self.name)
        self.attributes = {}
        self.other_attributes = {}
        self.possible_attributes = {
            'landscape' : ('False', ['True', 'False']),
            'page-border' : ('none', ['none', 'single', 'single-thick',
                                     'double', 'double-thick']),
            }

        for key, value in attrs.items():
            if key.endswith("-default"):
                name = key[:-len("-default")]
                if name in ["job-sheets", "printer-error-policy",
                            "printer-op-policy", # handled below
                            "notify-events", # cannot be set
                            "document-format", # cannot be set
                            "notify-lease-duration"]: # cannot be set
                    continue

                supported = attrs.get(name + "-supported", None) or \
                            self.possible_attributes.get(name, None) or \
                            ""

                # Convert a list into a comma-separated string, since
                # it can only really have been misinterpreted as a list
                # by CUPS.
                if isinstance (value, list):
                    value = reduce (lambda x, y: x+','+y, value)

                self.attributes[name] = value

                if name+"-supported" in attrs:
                    supported = attrs[name+"-supported"]
                    self.possible_attributes[name] = (value, supported)
            elif (not key.endswith ("-supported") and
                  key != 'job-sheets-default' and
                  key != 'printer-error-policy' and
                  key != 'printer-op-policy' and
                  not key.startswith ('requesting-user-name-')):
                self.other_attributes[key] = value

        self.job_sheet_start, self.job_sheet_end = attrs.get(
            'job-sheets-default', ('none', 'none'))
        self.job_sheets_supported = attrs.get('job-sheets-supported', ['none'])
        self.error_policy = attrs.get('printer-error-policy', 'none')
        self.error_policy_supported = attrs.get(
            'printer-error-policy-supported', ['none'])
        self.op_policy = attrs.get('printer-op-policy', "") or "default"
        self.op_policy_supported = attrs.get(
            'printer-op-policy-supported', ["default"])

        self.default_allow = True
        self.except_users = []
        if 'requesting-user-name-allowed' in attrs:
            self.except_users = attrs['requesting-user-name-allowed']
            self.default_allow = False
        elif 'requesting-user-name-denied' in attrs:
            self.except_users = attrs['requesting-user-name-denied']
        self.except_users_string = ', '.join(self.except_users)
        self.update (**attrs)

    def getServer(self):
        """
        Find out which server defines this printer.

        @returns: server URI or None
        """
        if not self.uri_supported[0].startswith('ipp://'):
            return None
        uri = self.uri_supported[0][6:]
        uri = uri.split('/')[0]
        uri = uri.split(':')[0]
        if uri == "localhost.localdomain":
            uri = "localhost"
        return uri

    def getPPD(self):
        """
        Obtain the printer's PPD.

        @returns: cups.PPD object, or False for raw queues
        @raise cups.IPPError: IPP error
        """
        result = None
        if self._ppd is None:
            try:
                self._ppd = self.connection.getPPD(self.name)
                result = cups.PPD (self._ppd)
            except cups.IPPError as emargs:
                (e, m) = emargs.args
                if e == cups.IPP_NOT_FOUND:
                    result = False
                else:
                    raise

        if result is None and self._ppd is not None:
            result = cups.PPD (self._ppd)

        return result

    def setOption(self, name, value):
        """
        Set a printer's option.

        @param name: option name
        @type name: string
        @param value: option value
        @type value: option-specific
        """
        if isinstance (value, float):
            radixchar = locale.nl_langinfo (locale.RADIXCHAR)
            if radixchar != '.':
                # Convert floats to strings, being careful with decimal points.
                value = str (value).replace (radixchar, '.')
        self.connection.addPrinterOptionDefault(self.name, name, value)

    def unsetOption(self, name):
        """
        Unset a printer's option.

        @param name: option name
        @type name: string
        """
        self.connection.deletePrinterOptionDefault(self.name, name)

    def setEnabled(self, on, reason=None):
        """
        Set the printer's enabled state.

        @param on: whether it will be enabled
        @type on: bool
        @param reason: reason for this state
        @type reason: string
        """
        if on:
            self.connection.enablePrinter(self.name)
        else:
            if reason:
                self.connection.disablePrinter(self.name, reason=reason)
            else:
                self.connection.disablePrinter(self.name)

    def setAccepting(self, on, reason=None):
        """
        Set the printer's accepting state.

        @param on: whether it will be accepting
        @type on: bool
        @param reason: reason for this state
        @type reason: string
        """
        if on:
            self.connection.acceptJobs(self.name)
        else:
            if reason:
                self.connection.rejectJobs(self.name, reason=reason)
            else:
                self.connection.rejectJobs(self.name)

    def setShared(self,on):
        """
        Set the printer's shared state.

        @param on: whether it will be accepting
        @type on: bool
        """
        self.connection.setPrinterShared(self.name, on)

    def setErrorPolicy (self, policy):
        """
        Set the printer's error policy.

        @param policy: error policy
        @type policy: string
        """
        self.connection.setPrinterErrorPolicy(self.name, policy)

    def setOperationPolicy(self, policy):
        """
        Set the printer's operation policy.

        @param policy: operation policy
        @type policy: string
        """
        self.connection.setPrinterOpPolicy(self.name, policy)

    def setJobSheets(self, start, end):
        """
        Set the printer's job sheets.

        @param start: start sheet
        @type start: string
        @param end: end sheet
        @type end: string
        """
        self.connection.setPrinterJobSheets(self.name, start, end)

    def setAccess(self, allow, except_users):
        """
        Set access control list.

        @param allow: whether to allow by default, otherwise deny
        @type allow: bool
        @param except_users: exception list
        @type except_users: string list
        """
        if isinstance(except_users, str):
            users = except_users.split()
            users = [u.split(",") for u in users]
            except_users = []
            for u in users:
                except_users.extend(u)
            except_users = [u.strip() for u in except_users]
            except_users = [_f for _f in except_users if _f]

        if allow:
            self.connection.setPrinterUsersDenied(self.name, except_users)
        else:
            self.connection.setPrinterUsersAllowed(self.name, except_users)

    def jobsQueued(self, only_tests=False, limit=None):
        """
        Find out whether jobs are queued for this printer.

        @param only_tests: whether to restrict search to test pages
        @type only_tests: bool
        @returns: list of job IDs
        """
        ret = []
        try:
            try:
                r = ['job-id', 'job-printer-uri', 'job-name']
                jobs = self.connection.getJobs (requested_attributes=r)
            except TypeError:
                # requested_attributes requires pycups 1.9.50
                jobs = self.connection.getJobs ()
        except cups.IPPError:
            return ret

        for id, attrs in jobs.items():
            try:
                uri = attrs['job-printer-uri']
                uri = uri[uri.rindex ('/') + 1:]
            except:
                continue
            if uri != self.name:
                continue

            if (not only_tests or
                ('job-name' in attrs and
                 attrs['job-name'] == 'Test Page')):
                ret.append (id)

                if limit is not None and len (ret) == limit:
                    break
        return ret

    def jobsPreserved(self, limit=None):
        """
        Find out whether there are preserved jobs for this printer.

        @return: list of job IDs
        """
        ret = []
        try:
            try:
                r = ['job-id', 'job-printer-uri', 'job-state']
                jobs = self.connection.getJobs (which_jobs='completed',
                                                requested_attributes=r)
            except TypeError:
                # requested_attributes requires pycups 1.9.50
                jobs = self.connection.getJobs (which_jobs='completed')
        except cups.IPPError:
            return ret

        for id, attrs in jobs.items():
            try:
                uri = attrs['job-printer-uri']
                uri = uri[uri.rindex ('/') + 1:]
            except:
                continue
            if uri != self.name:
                continue
            if (attrs.get ('job-state',
                           cups.IPP_JOB_PENDING) < cups.IPP_JOB_COMPLETED):
                continue
            ret.append (id)
            if limit is not None and len (ret) == limit:
                break

        return ret

    def testsQueued(self, limit=None):
        """
        Find out whether test jobs are queued for this printer.

        @returns: list of job IDs
        """
        return self.jobsQueued (only_tests=True, limit=limit)

    def setAsDefault(self):
        """
        Set this printer as the system default.
        """
        self.connection.setDefault(self.name)

        # Also need to check system-wide lpoptions because that's how
        # previous Fedora versions set the default (bug #217395).
        with tempfile.TemporaryFile () as f:
            try:
                resource = "/admin/conf/lpoptions"
                self.connection.getFile(resource, fd=f.fileno ())
            except cups.HTTPError as e:
                (s,) = e.args
                if s in [cups.HTTP_NOT_FOUND, cups.HTTP_AUTHORIZATION_CANCELED]:
                    return False

                raise cups.HTTPError (s)

            f.seek (0)
            lines = [ line.decode('UTF-8') for line in f.readlines () ]
            changed = False
            i = 0
            for line in lines:
                if line.startswith ("Default "):
                    # This is the system-wide default.
                    name = line.split (' ')[1]
                    if name != self.name:
                        # Stop it from over-riding the server default.
                        lines[i] = "Dest " + line[8:]
                        changed = True
                    i += 1

            if changed:
                f.seek (0)
                f.writelines ([ line.encode('UTF-8') for line in lines ])
                f.truncate ()
                f.flush ()
                f.seek (0)
                try:
                    self.connection.putFile (resource, fd=f.fileno ())
                except cups.HTTPError:
                    return False

            return changed

def getPrinters(connection):
    """
    Obtain a list of printers.

    @param connection: CUPS connection
    @type connection: CUPS.Connection object
    @returns: L{Printer} list
    """
    printers = connection.getPrinters()
    classes = connection.getClasses()
    for name, printer in printers.items():
        printer = Printer(name, connection, **printer)
        printers[name] = printer
        if name in classes:
            printer.class_members = classes[name]
            printer.class_members.sort()
    return printers

def parseDeviceID (id):
    """
    Parse an IEEE 1284 Device ID, so that it may be indexed by field name.

    @param id: IEEE 1284 Device ID, without the two leading length bytes
    @type id: string
    @returns: dict indexed by field name
    """
    id_dict = {}
    pieces = id.split(";")
    for piece in pieces:
        if piece.find(":") == -1:
            continue
        name, value = piece.split(":",1)
        id_dict[name.strip ()] = value.strip()
    if "MANUFACTURER" in id_dict:
        id_dict.setdefault("MFG", id_dict["MANUFACTURER"])
    if "MODEL" in id_dict:
        id_dict.setdefault("MDL", id_dict["MODEL"])
    if "COMMAND SET" in id_dict:
        id_dict.setdefault("CMD", id_dict["COMMAND SET"])
    for name in ["MFG", "MDL", "CMD", "CLS", "DES", "SN", "S", "P", "J"]:
        id_dict.setdefault(name, "")
    if id_dict["CMD"] == '':
        id_dict["CMD"] = []
    else:
        id_dict["CMD"] = id_dict["CMD"].split(',')
    return id_dict

class Device:
    """
    This class represents a CUPS device.
    """

    def __init__(self, uri, **kw):
        """
        @param uri: device URI
        @type uri: string
        @param kw: device attributes
        @type kw: dict
        """
        self.uri = uri
        self.device_class = kw.get('device-class', '')
        self.info = kw.get('device-info', '')
        self.make_and_model = kw.get('device-make-and-model', '')
        self.id = kw.get('device-id', '')
        self.location = kw.get('device-location', '')

        uri_pieces = uri.split(":")
        self.type =  uri_pieces[0]
        self.is_class = len(uri_pieces)==1

        #self.id = 'MFG:HEWLETT-PACKARD;MDL:DESKJET 990C;CMD:MLC,PCL,PML;CLS:PRINTER;DES:Hewlett-Packard DeskJet 990C;SN:US05N1J00XLG;S:00808880800010032C1000000C2000000;P:0800,FL,B0;J:                    ;'

        self.id_dict = parseDeviceID (self.id)

        s = uri.find("serial=")
        if s != -1 and not self.id_dict.get ('SN',''):
            self.id_dict['SN'] = uri[s + 7:]

    def __repr__ (self):
        return "<cupshelpers.Device \"%s\">" % self.uri

    def __lt__(self, other):
        """
        Compare devices by order of preference.
        """
        if other is None:
            return True

        if self.is_class != other.is_class:
            if other.is_class:
                return True
            return False

        stype = self.type
        if stype == "dnssd":
            if self.uri.find("._ipp") != -1:
                stype = "dnssdi"
            elif self.uri.find("._pdl-datastream") != -1:
                stype = "dnssds"
            elif self.uri.find("._printer") != -1:
                stype = "dnssdl"
        if stype == "usb":
            if self.uri.lower().find("fax") != -1:
                stype = "usbfax"
        otype = other.type
        if otype == "dnssd":
            if other.uri.find("._ipp") != -1:
                otype = "dnssdi"
            elif other.uri.find("._pdl-datastream") != -1:
                otype = "dnssds"
            elif other.uri.find("._printer") != -1:
                otype = "dnssdl"
        if otype == "usb":
            if other.uri.lower().find("fax") != -1:
                otype = "usbfax"

        if not self.is_class and (stype != otype):
            # "hp"/"hpfax" before "usb" before * before "parallel" before
            # "serial"
            if otype == "serial":
                return True
            if stype == "serial":
                return False
            if otype == "parallel":
                return True
            if stype == "parallel":
                return False
            if otype == "hp":
                return False
            if stype == "hp":
                return True
            if otype == "hpfax":
                return False
            if stype == "hpfax":
                return True
            if otype == "dnssdi":
                return False
            if stype == "dnssdi":
                return True
            if otype == "ipps":
                return False
            if stype == "ipps":
                return True
            if otype == "ipp":
                return False
            if stype == "ipp":
                return True
            if otype == "dnssds":
                return False
            if stype == "dnssds":
                return True
            if otype == "socket":
                return False
            if stype == "socket":
                return True
            if otype == "dnssdl":
                return False
            if stype == "dnssdl":
                return True
            if otype == "lpd":
                return False
            if stype == "lpd":
                return True
            if otype == "usb":
                return False
            if stype == "usb":
                return True
            if otype == "usbfax":
                return False
            if stype == "usbfax":
                return True
        result = bool(self.id) < bool(other.id)
        if not result:
            result = self.info < other.info

        return result

class _GetDevicesCall(object):
    def call (self, connection, kwds):
        if "reply_handler" in kwds:
            self._client_reply_handler = kwds.get ("reply_handler")
            kwds["reply_handler"] = self._reply_handler
            return connection.getDevices (**kwds)

        self._client_reply_handler = None
        result = connection.getDevices (**kwds)
        return self._reply_handler (connection, result)

    def _reply_handler (self, connection, devices):
        for uri, data in devices.items():
            device = Device(uri, **data)
            devices[uri] = device
            if device.info != '' and device.make_and_model == '':
                device.make_and_model = device.info

        if self._client_reply_handler:
            self._client_reply_handler (connection, devices)
        else:
            return devices

def getDevices(connection, **kw):
    """
    Obtain a list of available CUPS devices.

    @param connection: CUPS connection
    @type connection: cups.Connection object
    @returns: a list of L{Device} objects
    @raise cups.IPPError: IPP Error
    """
    op = _GetDevicesCall ()
    return op.call (connection, kw)

def activateNewPrinter(connection, name):
    """
    Set a new printer enabled, accepting jobs, and (if necessary) the
    default printer.

    @param connection: CUPS connection
    @type connection: cups.Connection object
    @param name: printer name
    @type name: string
    @raise cups.IPPError: IPP error
    """
    connection.enablePrinter (name)
    connection.acceptJobs (name)

    # Set as the default if there is not already a default printer.
    if connection.getDefault () is None:
        connection.setDefault (name)

def copyPPDOptions(ppd1, ppd2):
    """
    Copy default options between PPDs.

    @param ppd1: source PPD
    @type ppd1: cups.PPD object
    @param ppd2: destination PPD
    @type ppd2: cups.PPD object
    """
    def getPPDGroupOptions(group):
        options = group.options[:]
        for g in group.subgroups:
            options.extend(getPPDGroupOptions(g))
        return options

    def iteratePPDOptions(ppd):
    	for group in ppd.optionGroups:
            for option in getPPDGroupOptions(group):
            	yield option

    for option in iteratePPDOptions(ppd1):
        if option.keyword == "PageRegion":
            continue
        new_option = ppd2.findOption(option.keyword)
        if new_option and option.ui==new_option.ui:
            value = option.defchoice
            for choice in new_option.choices:
                if choice["choice"]==value:
                    ppd2.markOption(new_option.keyword, value)
                    _debugprint ("set %s = %s" % (repr (new_option.keyword),
                                                  repr (value)))

def setPPDPageSize(ppd, language):
    """
    Set the PPD page size according to locale.

    @param ppd: PPD
    @type ppd: cups.PPD object
    @param language: language, as given by the first element of
    locale.setlocale
    @type language: string
    """
    # Just set the page size to A4 or Letter, that's all.
    # Use the same method CUPS uses.
    size = 'A4'
    letter = [ 'C', 'POSIX', 'en', 'en_US', 'en_CA', 'fr_CA' ]
    for each in letter:
        if language == each:
            size = 'Letter'
    try:
        ppd.markOption ('PageSize', size)
        _debugprint ("set PageSize = %s" % size)
    except:
        _debugprint ("Failed to set PageSize (%s not available?)" % size)

def missingExecutables(ppd):
    """
    Check that all relevant executables for a PPD are installed.

    @param ppd: PPD
    @type ppd: cups.PPD object
    @returns: string list, representing missing executables
    """

    # First, a local function.  How to check that something exists
    # in a path:
    def pathcheck (name, path="/usr/bin:/bin"):
        if name == "-":
            # A filter of "-" means that no filter is required,
            # i.e. the device accepts the given format as-is.
            return "builtin"
        # Strip out foomatic '%'-style place-holders.
        p = name.find ('%')
        if p != -1:
            name = name[:p]
        if len (name) == 0:
            return "true"
        if name[0] == '/':
            if os.access (name, os.X_OK):
                _debugprint ("%s: found" % name)
                return name
            else:
                _debugprint ("%s: NOT found" % name)
                return None
        if name.find ("=") != -1:
            return "builtin"
        if name in [ ":", ".", "[", "alias", "bind", "break", "cd",
                     "continue", "declare", "echo", "else", "eval",
                     "exec", "exit", "export", "fi", "if", "kill", "let",
                     "local", "popd", "printf", "pushd", "pwd", "read",
                     "readonly", "set", "shift", "shopt", "source",
                     "test", "then", "trap", "type", "ulimit", "umask",
                     "unalias", "unset", "wait" ]:
            return "builtin"
        for component in path.split (':'):
            file = component.rstrip (os.path.sep) + os.path.sep + name
            if os.access (file, os.X_OK):
                _debugprint ("%s: found" % file)
                return file
        _debugprint ("%s: NOT found in %s" % (name, path))
        return None

    exes_to_install = []

    def add_missing (exe):
        # Strip out foomatic '%'-style place-holders.
        p = exe.find ('%')
        if p != -1:
            exe = exe[:p]

        exes_to_install.append (exe)

    # Find a 'FoomaticRIPCommandLine' attribute.
    exe = exepath = None
    attr = ppd.findAttr ('FoomaticRIPCommandLine')
    if attr:
        # Foomatic RIP command line to check.
        cmdline = attr.value.replace ('&&\n', '')
        cmdline = cmdline.replace ('"', '"')
        cmdline = cmdline.replace ('<', '<')
        cmdline = cmdline.replace ('>', '>')
        if (cmdline.find ("(") != -1 or
            cmdline.find ("&") != -1):
            # Don't try to handle sub-shells or unreplaced HTML entities.
            cmdline = ""

        # Strip out foomatic '%'-style place-holders
        pipes = cmdline.split (';')
        for pipe in pipes:
            cmds = pipe.strip ().split ('|')
            for cmd in cmds:
                args = cmd.strip ().split (' ')
                exe = args[0]
                exepath = pathcheck (exe)
                if not exepath:
                    add_missing (exe)
                    continue

                # Main executable found.  But if it's 'gs',
                # perhaps there is an IJS server we also need
                # to check.
                if os.path.basename (exepath) == 'gs':
                    argn = len (args)
                    argi = 1
                    search = "-sIjsServer="
                    while argi < argn:
                        arg = args[argi]
                        if arg.startswith (search):
                            exe = arg[len (search):]
                            exepath = pathcheck (exe)
                            if not exepath:
                                add_missing (exe)

                            break

                        argi += 1

            if not exepath:
                # Next pipe.
                break

    if exepath or not exe:
        # Look for '*cupsFilter' lines in the PPD and check that
        # the filters are installed.
        (tmpfd, tmpfname) = tempfile.mkstemp (text=True)
        os.unlink (tmpfname)
        ppd.writeFd (tmpfd)
        os.lseek (tmpfd, 0, os.SEEK_SET)
        f = os.fdopen (tmpfd, "rt", encoding="utf-8")
        search = "*cupsFilter:"
        for line in f:
            if line.startswith (search):
                line = line[len (search):].strip ().strip ('"')
                try:
                    (mimetype, cost, exe) = line.split (' ')
                except:
                    continue

                exepath = pathcheck (exe,
                                     config.cupsserverbindir + "/filter:"
                                     "/usr/lib64/cups/filter")
                if not exepath:
                    add_missing (config.cupsserverbindir + "/filter/" + exe)

    return exes_to_install

def missingPackagesAndExecutables(ppd):
    """
    Check that all relevant executables for a PPD are installed.

    @param ppd: PPD
    @type ppd: cups.PPD object
    @returns: string list pair, representing missing packages and
    missing executables
    """
    executables = missingExecutables(ppd)
    return ([], executables)

def _main():
    c = cups.Connection()
    #printers = getPrinters(c)
    for device in getDevices(c).values():
        print (device.uri, device.id_dict)

if __name__=="__main__":
    _main()


import turtle
import time

def printIntro():
    print("该程序利用turtle完成对七段数码管日期显示的配置")
    print("程序运行需要笔刷的颜色")

def getInputs():
    while True:
        try:
            a=input("请输入笔刷的颜色(red blue,green)：")
            b=eval(input("请输入运笔速度(10-50)："))
            return a,b
        except Exception:
            print('输入错误请重新输入')

def drawGap():
    turtle.penup()
    turtle.fd(5)

def drawLine(draw):
    drawGap()
    turtle.pensize(8) if draw else turtle.pensize(1)
    turtle.pendown()
    turtle.left(45)
    turtle.fd(5)
    turtle.right(45)
    turtle.fd(40)
    turtle.right(45)
    turtle.fd(5)
    turtle.right(90)
    turtle.fd(5)
    turtle.right(45)
    turtle.fd(40)
    turtle.right(45)
    turtle.fd(5)
    turtle.right(135)
    turtle.penup()
    turtle.fd(47)
    turtle.pendown()
    drawGap()

def drawDigit(digit):

    drawLine(True) if digit in [2,3,4,5,6,8,9] else drawLine(False)
    turtle.right(100)
    drawLine(True) if digit in [0,1,3,4,5,6,7,8,9] else drawLine(False)
    turtle.right(80)
    drawLine(True) if digit in [0,2,3,5,6,8,9] else drawLine(False)
    turtle.right(100)
    drawLine(True) if digit in [0,2,6,8] else drawLine(False)

    drawLine(True) if digit in [0,4,5,6,8,9] else drawLine(False)
    turtle.right(80)
    drawLine(True) if digit in [0,2,3,5,6,7,8,9] else drawLine(False)
    turtle.right(100)
    drawLine(True) if digit in [0,1,2,3,4,7,8,9] else drawLine(False)
    turtle.right(80)
    turtle.left(180)
    turtle.penup()
    turtle.fd(20)

def drawDate(date):
    try:
        for i in date:
            if(i=='年'):
                turtle.write('年',font=("Arial",18,"normal"))
                turtle.pencolor("red")
                turtle.fd(40)
            elif i=='月':
                turtle.write('月',font=("Arial",18,"normal"))
                turtle.fd(40)
            elif i=='日':
                turtle.write('日',font=("Arial",18,"normal"))
                turtle.goto(-300,0)
                turtle.right(90)
                turtle.fd(150)
                turtle.left(90)

            else:
                drawDigit(eval(i))
    except NameError:
        print("输入格式错误")

def draw(color,speed):
    turtle.setup(1000,600,0,0)
    turtle.pencolor(color)
    turtle.speed(speed)
    turtle.hideturtle()
    turtle.penup()
    turtle.fd(-300)
    drawDate(time.strftime('%Y年%m月%d日',time.gmtime()))
    drawDate('abc')
    turtle.done()

def printSummary():
    print("已成功执行上述程序")

def main():
    printIntro()
    color, speed=getInputs()
    draw(color,speed)
    printSummary()


import numpy as np
from bathygrid.algorithms import *
from bathygrid.maingrid import *
import matplotlib.pyplot as plt

#
# x = np.array([200, 250, 300, 350, 400, 450, 500, 550, 600, 650, 700, 750, 800, 400, 600])
# y = np.array([600, 550, 500, 470, 440, 410, 400, 410, 440, 470, 500, 550, 600, 800, 800])
# x = x.ravel()
# y = y.ravel()
# z = np.linspace(5, 30, num=x.size).astype(np.float32)
# tvu = np.linspace(1, 2, num=x.size).astype(np.float32)
# thu = np.linspace(0.5, 1, num=x.size).astype(np.float32)
#
# dtyp = [('x', np.float64), ('y', np.float64), ('z', np.float32), ('tvu', np.float32), ('thu', np.float32)]
# smileyface = np.empty(len(x), dtype=dtyp)
# smileyface['x'] = x
# smileyface['y'] = y
# smileyface['z'] = z
# smileyface['tvu'] = tvu
# smileyface['thu'] = thu
#
# bg = SRGrid(tile_size=1024)
# bg.add_points(smileyface, 'test1', ['line1', 'line2'], 26917, 'waterline')
# bg.grid(resolution=64)
# plt.scatter(smileyface['x'], smileyface['y'], c=smileyface['z'])
# bg.plot()
#
# plt.figure()
# x, y, lyrdata, newmins, newmaxs = bg.return_surf_xyz('depth', 64.0, False)
# lat2d, lon2d = np.meshgrid(y, x)
# data_m = np.ma.array(lyrdata[0], mask=np.isnan(lyrdata[0]))
# plt.pcolormesh(lon2d, lat2d, data_m.T, vmin=data_m.min(), vmax=data_m.max())

def trial_data():
    number_of_points = 100000
    depth = np.linspace(10, 20, number_of_points)
    vunc = np.linspace(1, 2, number_of_points)
    hunc = np.linspace(1, 2, number_of_points)
    cell_indices = np.random.randint(0, 10000, number_of_points)
    densitygrid = np.full((100, 100), 0)
    grid = np.full((100, 100), np.nan)
    vuncgrid = np.full((100, 100), np.nan)
    huncgrid = np.full((100, 100), np.nan)
    return depth, vunc, hunc, cell_indices, grid, densitygrid, vuncgrid, huncgrid

# def trial_grid_mean_numba():
#     depth, tvu, thu, cell_indices, grid, tvugrid, thugrid = trial_data()
#     nb_grid_mean(depth, cell_indices, grid, tvu, thu, tvugrid, thugrid)

def trial_grid_mean_numpy():
    depth, tvu, thu, cell_indices, grid, densitygrid, tvugrid, thugrid = trial_data()
    np_grid_mean(depth, cell_indices, grid, densitygrid, tvu, thu, tvugrid, thugrid)

# def trial_grid_shoal_numba():
#     depth, tvu, thu, cell_indices, grid, tvugrid, thugrid = trial_data()
#     nb_grid_shoalest(depth, cell_indices, grid, tvu, thu, tvugrid, thugrid)

def trial_grid_shoal_numpy():
    depth, tvu, thu, cell_indices, grid, densitygrid, tvugrid, thugrid = trial_data()
    np_grid_shoalest(depth, cell_indices, grid, densitygrid, tvu, thu, tvugrid, thugrid)

if __name__ == '__main__':
    import timeit

    # trial_grid_mean_numba()  # run once to skip the overhead around import/compiling
    # trial_grid_shoal_numba()  # run once to skip the overhead around import/compiling

    # print('Numba mean: {}'.format(timeit.timeit(trial_grid_mean_numba, number=2)))
    print('Numpy mean: {}'.format(timeit.timeit(trial_grid_mean_numpy, number=2)))
    # print('Numba shoal: {}'.format(timeit.timeit(trial_grid_shoal_numba, number=2)))
    print('Numpy shoal: {}'.format(timeit.timeit(trial_grid_shoal_numpy, number=2)))