"""
Metrics used for evaluating and comparing partial labels models

Heavily based on:
* https://github.com/IndicoDataSolutions/finetune/blob/development/finetune/util/metrics.py
* https://github.com/IndicoDataSolutions/benchner/blob/rd/synthetic/benchner/metrics.py

For all metric calculations below:
- A "label" or "prediction" is a dictionary of the following format:
    {"start": int, "end": int, "text": str, "label": str, "doc_idx": int}
- "true" or "predicted" generally refers to a list of label dictionaries of the format above
"""
import traceback
from collections import defaultdict, OrderedDict
import copy
import functools
import random
from typing import Callable, Dict, List, Set, Tuple, Union

# Need extra packages to use these metrics
# TODO Figure out how to incorporate spacy model download into package install
try:
    import numpy as np
    import spacy
except ImportError:
    traceback.print_exc()
    raise Exception("Using partial labels requires installing additional packages")

@functools.lru_cache()
def get_spacy():
    """
    FIXME return type of spacy.lang.en.English not working for type annotation

    Load English core language model from spacy to use for tokenization. Wrapped
    with lru_cache to ensure that the model is only loaded once.

    Returns:
        nlp: Spacy model
    """
    nlp = spacy.load("en_core_web_sm", disable=["parser", "tagger", "ner", "textcat"])
    nlp.max_length = (
        800000000  # approximately one volume of the encyclopedia britannica.
    )
    return nlp

def calc_recall(num_tps: int, num_fns: int) -> float:
    """
    Calculate recall for a class given the number of true positives and false negatives

    Args:
        num_tps: Number of true positives
        num_fns: Number of false negatives

    Returns:
        recall
    """
    try:
        return num_tps / float(num_fns + num_tps)
    except ZeroDivisionError:
        return 0.0

def calc_precision(num_tps: int, num_fps: int) -> float:
    """
    Calculate precision for a class given the number of true positives and false positives

    Args:
        num_tps: Number of true positives
        num_fps: Number of false positives

    Returns:
        precision
    """
    try:
        return num_tps / float(num_fps + num_tps)
    except ZeroDivisionError:
        return 0.0

def calc_f1(recall: float, precision: float) -> float:
    """
    Calculate F1 for a class given the recall and precision

    Args:
        recall
        precision

    Returns:
        f1
    """
    try:
        return 2 * (recall * precision) / (recall + precision)
    except ZeroDivisionError:
        return 0.0

def _get_unique_classes(
    true: List[List[Dict[str, Union[int, str]]]],
    predicted: List[List[Dict[str, Union[int, str]]]],
) -> List[str]:
    """
    Given ground truth labels and predictions, take the union of label classes
    to determine the set of unique classes

    Args:
        true: Ground truth labels for a list of documents
        predicted: Model predictions for a list of documents

    Returns:
        unique_classes: A list of unique classes
    """
    true_and_pred = list(true) + list(predicted)
    return list(set([seq["label"] for seqs in true_and_pred for seq in seqs]))

def _convert_to_token_list(
    annotations: List[Dict[str, Union[int, str]]], doc_idx: int = None
) -> List[Dict[str, Union[int, str]]]:
    """
    Given annotations (labels or predictions) for a single document, use the spacy
    English tokenizer to split the annotations into individual tokens

    Args:
        annotations: Labels or predictions for a single document
        doc_idx: Document index for current document

    Returns:
        tokens
    """
    nlp = get_spacy()
    tokens = []
    annotations = copy.deepcopy(annotations)

    for annotation in annotations:
        start_idx = annotation.get("start")
        tokens.extend(
            [
                {
                    "start": start_idx + token.idx,
                    "end": start_idx + token.idx + len(token.text),
                    "text": token.text,
                    "label": annotation.get("label"),
                    "doc_idx": doc_idx,
                }
                for token in nlp(annotation.get("text"))
            ]
        )

    return tokens

def sequence_labeling_token_conf_matrix(
    true: List[List[Dict[str, Union[int, str]]]],
    predicted: List[List[Dict[str, Union[int, str]]]],
) -> Dict[str, Dict[str, List[Dict[str, Union[int, str]]]]]:
    """
    Given ground truth labels and predictions, determine the true positives (TPs),
    false positives (FPs), and false negatives (FNs) for each class at a per token level.

    Args:
        true: Ground truth labels for a list of documents
        predicted: Model predictions for a list of documents

    Returns:
        token_conf_matrix: A dictionary containing TP/FP/FN counts for each class:
            {"classA": {"true_positives": List[Dict], "false_positives": List[Dict],
                        "false_negatives": List[Dict]}, ...}
    """
    # Determine set of unique classes, and set up token_count dict for each class
    unique_classes = _get_unique_classes(true, predicted)
    token_conf_matrix = {
        cls: {"true_positives": [], "false_positives": [], "false_negatives": []}
        for cls in unique_classes
    }

    for i, (true_list, pred_list) in enumerate(zip(true, predicted)):
        # Convert span annotations to token annotations
        true_tokens = _convert_to_token_list(true_list, doc_idx=i)
        pred_tokens = _convert_to_token_list(pred_list, doc_idx=i)

        # Calculate true positives and false negatives
        for true_token in true_tokens:
            for pred_token in pred_tokens:
                if (
                    pred_token["start"] == true_token["start"]
                    and pred_token["end"] == true_token["end"]
                ):

                    if pred_token["label"] == true_token["label"]:
                        token_conf_matrix[true_token["label"]]["true_positives"].append(
                            true_token
                        )
                    else:
                        token_conf_matrix[true_token["label"]][
                            "false_negatives"
                        ].append(true_token)
                        token_conf_matrix[pred_token["label"]][
                            "false_positives"
                        ].append(pred_token)

                    break
            else:
                token_conf_matrix[true_token["label"]]["false_negatives"].append(
                    true_token
                )

        # Calculate false positives
        for pred_token in pred_tokens:
            for true_token in true_tokens:
                if (
                    pred_token["start"] == true_token["start"]
                    and pred_token["end"] == true_token["end"]
                ):
                    break
            else:
                token_conf_matrix[pred_token["label"]]["false_positives"].append(
                    pred_token
                )

    return token_conf_matrix

def sequence_labeling_conf_matrix(
    true: List[List[Dict[str, Union[int, str]]]],
    predicted: List[List[Dict[str, Union[int, str]]]],
    equality_fn: Callable[
        [Dict[str, Union[int, str]], Dict[str, Union[int, str]]], bool
    ],
) -> Dict[str, Dict[str, List[Dict[str, Union[int, str]]]]]:
    """
    Given ground truth labels and predictions, determine the true positives (TPs),
    false positives (FPs), and false negatives (FNs) for each class.

    The method to determine TPs/FPs/FNs is based on the equality function,
    and can be exact match, overlap, or superset

    Args:
        true: Ground truth labels for a list of documents
        predicted: Model predictions for a list of documents
        equality_fn: Function to determine equality between two spans

    Returns:
        conf_matrix: A dictionary containing TP/FP/FN counts for each class:
            {"classA": {"true_positives": List[Dict], "false_positives": List[Dict],
                        "false_negatives": List[Dict]}, ...}
    """
    # Determine set of unique classes, and set up count dict for each class
    unique_classes = _get_unique_classes(true, predicted)
    conf_matrix = {
        cls: {"true_positives": [], "false_positives": [], "false_negatives": []}
        for cls in unique_classes
    }

    for i, (true_annotations, predicted_annotations) in enumerate(zip(true, predicted)):
        # Add doc idx to make verification easier
        for annotations in [true_annotations, predicted_annotations]:
            for annotation in annotations:
                annotation["doc_idx"] = i

        # Calculate true positives and false negatives
        for true_annotation in true_annotations:
            for pred_annotation in predicted_annotations:
                # Use equality function to determine match between ground truth and prediction
                if equality_fn(true_annotation, pred_annotation):
                    if pred_annotation["label"] == true_annotation["label"]:
                        conf_matrix[true_annotation["label"]]["true_positives"].append(
                            true_annotation
                        )
                        break
            else:
                conf_matrix[true_annotation["label"]]["false_negatives"].append(
                    true_annotation
                )

        # Calculate false positives
        for pred_annotation in predicted_annotations:
            for true_annotation in true_annotations:
                if (
                    equality_fn(true_annotation, pred_annotation)
                    and true_annotation["label"] == pred_annotation["label"]
                ):
                    break
            else:
                conf_matrix[pred_annotation["label"]]["false_positives"].append(
                    pred_annotation
                )

    return conf_matrix

def sequences_overlap(
    true_seq: Dict[str, Union[int, str]], pred_seq: Dict[str, Union[int, str]]
) -> bool:
    """
    Determine if two annotations overlap

    Args:
        true_seq: Ground truth label
        pred_seq: Prediction

    Both true_seq and pred_seq are label dictionaries

    Returns:
        is_match: True if sequences overlap
    """
    return true_seq["start"] < pred_seq["end"] and pred_seq["start"] < true_seq["end"]

def sequence_exact_match(
    true_seq: Dict[str, Union[int, str]], pred_seq: Dict[str, Union[int, str]]
) -> bool:
    """
    Determine if two annotations are an exact match (with stripped whitespace)

    Args:
        true_seq: Ground truth label
        pred_seq: Prediction

    Both true_seq and pred_seq are label dictionaries

    Returns:
        is_match: True if sequences are an exact match
    """
    true_seq = strip_whitespace(true_seq)
    pred_seq = strip_whitespace(pred_seq)
    return pred_seq["start"] == true_seq["start"] and pred_seq["end"] == true_seq["end"]

def sequence_superset(
    true_seq: Dict[str, Union[int, str]], pred_seq: Dict[str, Union[int, str]]
) -> bool:
    """
    Given a ground truth label and a prediction, determine if the prediction is
    a superset of the ground truth label (with stripped whitespace)

    Args:
        true_seq: Ground truth label
        pred_seq: Prediction

    Both true_seq and pred_seq are label dictionaries

    Returns:
        is_match: True if pred_seq is a superset of true_seq
    """
    true_seq = strip_whitespace(true_seq)
    pred_seq = strip_whitespace(pred_seq)
    return pred_seq["start"] <= true_seq["start"] and pred_seq["end"] >= true_seq["end"]

def get_seq_conf_matrix_fn(
    span_type: str = "token",
) -> Callable[
    [List[List[Dict[str, Union[int, str]]]], List[List[Dict[str, Union[int, str]]]]],
    Dict[str, Dict[str, List[Dict[str, Union[int, str]]]]],
]:
    """
    Determine what function to use for determining TPs/FPs/FNs based on the
    type of equality function to use.

    Args:
        span_type: Type of "span" to use to determine the equality function
            to use. Must be one of {"token", "equality", "partial", "overlap"}

    Returns:
        seq_conf_matrix_fn: Function to determine TPs/FPs/FNs
    """
    valid_span_types = {"token", "equality", "partial", "overlap"}
    if span_type not in valid_span_types:
        raise ValueError(
            f"span_type={span_type} is not valid, must be within {valid_span_types}"
        )

    span_type_fn_mapping = {
        "token": sequence_labeling_token_conf_matrix,
        "overlap": functools.partial(
            sequence_labeling_conf_matrix, equality_fn=sequences_overlap
        ),
        "exact": functools.partial(
            sequence_labeling_conf_matrix, equality_fn=sequence_exact_match
        ),
        "superset": functools.partial(
            sequence_labeling_conf_matrix, equality_fn=sequence_superset
        ),
    }
    return span_type_fn_mapping[span_type]

def strip_whitespace(
    annotation: Dict[str, Union[int, str]]
) -> Dict[str, Union[int, str]]:
    """
    Given a label annotation, strip whitespace from the start and end of the
    string and adjust the start/end indices of the annotation

    Args:
        annotation: Label annotation

    Returns:
        new_annotation: Updated annotation with stripped whitespace
    """
    label_text = annotation["text"]
    lstripped = label_text.lstrip()
    new_start = annotation["start"] + (len(label_text) - len(lstripped))
    stripped = label_text.strip()
    return {
        "text": label_text.strip(),
        "start": new_start,
        "end": new_start + len(stripped),
        "label": annotation["label"],
    }

def seq_recall(
    true: List[List[Dict[str, Union[int, str]]]],
    predicted: List[List[Dict[str, Union[int, str]]]],
    span_type: str = "token",
) -> Dict[str, float]:
    """
    Given ground truth labels and predictions, determine the recall for
    each class

    The method to determine TPs/FPs/FNs is based on the span type, which
    determines the equality function to compare ground truth and preds

    Args:
        true: Ground truth labels for a list of documents
        predicted: Model predictions for a list of documents
        span_type: Type of "span" to use to determine the equality function
            to use. Must be one of {"token", "equality", "partial", "overlap"}

    Returns:
        recalls: Recall for each class
    """
    valid_span_types = {"token", "equality", "partial", "overlap"}
    if span_type not in valid_span_types:
        raise ValueError(
            f"span_type={span_type} is not valid, must be within {valid_span_types}"
        )

    conf_matrix_fn = get_seq_conf_matrix_fn(span_type)
    # Get counts of TP/FP/FN for each class
    class_conf_matrices = conf_matrix_fn(true, predicted)
    recalls = {}
    for cls, conf_matrix in class_conf_matrices.items():
        num_tps = len(conf_matrix["true_positives"])
        num_fns = len(conf_matrix["false_negatives"])
        recalls[cls] = calc_recall(num_tps, num_fns)
    return recalls

def seq_precision(
    true: List[List[Dict[str, Union[int, str]]]],
    predicted: List[List[Dict[str, Union[int, str]]]],
    span_type: str = "token",
) -> Dict[str, float]:
    """
    Given ground truth labels and predictions, determine the precision for
    each class

    The method to determine TPs/FPs/FNs is based on the span type, which
    determines the equality function to compare ground truth and preds

    Args:
        true: Ground truth labels for a list of documents
        predicted: Model predictions for a list of documents
        span_type: Type of "span" to use to determine the equality function
            to use. Must be one of {"token", "equality", "partial", "overlap"}

    Returns:
        precisions: Precision for each class
    """
    valid_span_types = {"token", "equality", "partial", "overlap"}
    if span_type not in valid_span_types:
        raise ValueError(
            f"span_type={span_type} is not valid, must be within {valid_span_types}"
        )

    conf_matrix_fn = get_seq_conf_matrix_fn(span_type)
    # Get counts of TP/FP/FN for each class
    class_conf_matrices = conf_matrix_fn(true, predicted)
    precisions = {}
    for cls, conf_matrix in class_conf_matrices.items():
        num_tps = len(conf_matrix["true_positives"])
        num_fps = len(conf_matrix["false_positives"])
        precisions[cls] = calc_recall(num_tps, num_fps)
    return precisions

def calc_micro_f1(
    true: List[List[Dict[str, Union[int, str]]]],
    predicted: List[List[Dict[str, Union[int, str]]]],
    span_type: str = "token",
    include_recall_and_precision: bool = False,
) -> Union[Tuple[float, float, float], float]:
    """
    Given ground truth labels and predictions, determine the micro F1 score.
    Micro means calculating metrics globally by counting the total true positives,
    false negatives and false positives.

    Args:
        true: Ground truth labels for a list of documents
        predicted: Model predictions for a list of documents
        span_type: Type of "span" to use to determine the equality function
            to use. Must be one of {"token", "equality", "partial", "overlap"}
        include_recall_and_precision: If True, return a Tuple including the micro
            recall, precision, and F1. Else just return F1

    Returns:
        If include_recall_and_precision is True:
            recall
            precision
            f1
        Else:
            f1
    """
    conf_matrix_fn = get_seq_conf_matrix_fn(span_type)
    class_conf_matrices = conf_matrix_fn(true, predicted)
    true_positives, false_positives, false_negatives = 0, 0, 0
    for _, conf_matrix in class_conf_matrices.items():
        false_negatives += len(conf_matrix["false_negatives"])
        true_positives += len(conf_matrix["true_positives"])
        false_positives += len(conf_matrix["false_positives"])
    recall = calc_recall(true_positives, false_negatives)
    precision = calc_precision(true_positives, false_positives)
    f1 = calc_f1(recall, precision)
    if include_recall_and_precision:
        return recall, precision, f1
    return f1

def calc_per_class_f1(
    true: List[List[Dict[str, Union[int, str]]]],
    predicted: List[List[Dict[str, Union[int, str]]]],
    span_type: str = "token",
) -> Dict[str, Dict[str, Union[float, int]]]:
    """
    Given ground truth labels and predictions, determine the support (TPs+FNs)
    and F1 score for each class.

    Args:
        true: Ground truth labels for a list of documents
        predicted: Model predictions for a list of documents
        span_type: Type of "span" to use to determine the equality function
            to use. Must be one of {"token", "equality", "partial", "overlap"}

    Returns:
        results: OrderedDict containing support (TPs + FNs) and F1 score per class
            Note: Using Dict in return val type annotation to prevent TypeError
    """
    conf_matrix_fn = get_seq_conf_matrix_fn(span_type)
    class_conf_matrices = conf_matrix_fn(true, predicted)
    results = OrderedDict()
    for cls, conf_matrix in class_conf_matrices.items():
        results[cls] = {}
        num_tps = len(conf_matrix["true_positives"])
        num_fps = len(conf_matrix["false_positives"])
        num_fns = len(conf_matrix["false_negatives"])
        # Calculate per class precision and recall, and then F1
        recall = calc_recall(num_tps, num_fns)
        precision = calc_precision(num_tps, num_fps)
        results[cls]["support"] = num_tps + num_fns
        results[cls]["f1_score"] = calc_f1(recall, precision)
    return results

def calc_sequence_f1(
    true: List[List[Dict[str, Union[int, str]]]],
    predicted: List[List[Dict[str, Union[int, str]]]],
    span_type: str = "token",
    average: str = None,
) -> Union[float, Dict[str, Dict[str, Union[float, int]]]]:
    """
    Given ground truth labels and predictions, determine F1 score. The F1
    score can be calculated in different ways depending on the average
    parameter, which could be "micro", "macro", or "weighted" to calculate
    a single float value.

    If average is None, instead calculate per class F1 scores

    If average = None, return per-class F1 scores, otherwise
    return the requested model-level score.

    Args:
        true: Ground truth labels for a list of documents
        predicted: Model predictions for a list of documents
        span_type: Type of "span" to use to determine the equality function
            to use. Must be one of {"token", "equality", "partial", "overlap"}
        average: Metric to use to calculate F1. This is explained further in sklearn:
            https://scikit-learn.org/stable/modules/generated/sklearn.metrics.f1_score.html

    Returns:
        f1: F1 score in float form if average is "micro", "macro", or "weighted".
            Otherwise an OrderedDict with per class F1 scores
    """
    if average == "micro":
        return calc_micro_f1(true, predicted, span_type)

    f1s_by_class = calc_per_class_f1(true, predicted, span_type)
    f1s = [value.get("f1_score") for key, value in f1s_by_class.items()]
    supports = [value.get("support") for key, value in f1s_by_class.items()]

    if average == "weighted":
        return np.average(np.array(f1s), weights=np.array(supports))
    if average == "macro":
        return np.average(f1s)
    else:
        return f1s_by_class

################################################################
# All metrics below are not based on existing finetune metrics #
################################################################

def find_overlapping_sequence(
    pred: Dict[str, Union[int, str]], labels: List[Dict[str, Union[int, str]]]
) -> Union[Dict[str, Union[int, str]], None]:
    """
    Given a prediction and a set of ground truth labels, find the label sequence
    that overlaps with the prediction.

    Return None if no sequence is found

    Args:
        pred: Individual prediction (label dictionary)
        labels: List of ground truth label dictionaries for a document

    Returns:
        overlapping_label
    """
    for label in labels:
        if sequences_overlap(pred, label):
            return label

    return None

def identify_tps_fps_fns(
    true: List[List[Dict[str, Union[int, str]]]],
    predicted: List[List[Dict[str, Union[int, str]]]],
    tp_ratio: float = 1,
    fp_ratio: float = 1,
    fn_ratio: float = 1,
    random_seed: int = 42,
    tp_match_type: str = "exact",
    fp_match_type: str = "overlap",
    fn_match_type: str = "exact",
    replace_fps: bool = False,
    missing_label_ratio: float = None,
) -> Tuple[
    List[List[Dict[str, Union[int, str]]]],
    List[List[Dict[str, Union[int, str]]]],
    List[List[Dict[str, Union[int, str]]]],
]:
    """
    Given ground truth labels and predictions for a sequence labeling model,
    we want to identify the true positives, false positives, and false negatives.
    If any of the ratios (tp_ratio/fp_ratio/fn_ratio) are left at the defaults (0),
    we return an empty list. If the ratio is a float between 0 and 1, we sample
    that ratio of TPs, FPs, or FNs.

    The sequence_labeling_counts already identifies true positives, false positives,
    and false negatives, but in a significantly different format from what is expected by finetune.
    The format coming out of sequence_labeling_counts() is as follows:
    {
        <LABEL>: {
            "true_positives": int,
            "false_positives": int,
            "false_negatives": int
        },
        ...
    }

    The format expected by finetune is as follows, where annotations are separated
    for each document, and NOT separated by the different labels:
    [
        # Document 0
        [
            {
                "start": int, "end": int, "text": str, "label": str, "doc_idx": 0,
                "type": OneOf("tp", "fp", "fn")
            },
            ...
        ],
        # Document 1
        [...],
        ...
    ]

    Args:
        true: List of lists, where each sublist contains ground truth annotations
            for a given document
        predicted: List of lists, where each sublist contains predictions for a given document
        tp_ratio: Ratio of true positives to include, sampled randomly across all documents
        fp_ratio: Ratio of false positives to include, sampled randomly across all documents
        fn_ratio: Ratio of false negatives to include, sampled randomly across all documents
        random_seed: Random seed to use when sampling
        tp_match_type: Type of match in finetune metrics to use for identifying true positives.
            Must be one of ("exact", "overlap")
        fp_match_type: Same as above, but for false positives
        fn_match_type: Same as above, but for false negatives
        replace_fps: If we replace false positives with overlapping ground truth
            Very strong supervision, because we're relying on the reviewer or labeler
            to tell us the class and start/end of the overlapping ground truth
        missing_label_ratio: Used for simulating missing label classes for experiments.
            A ratio of X means that we will keep (1 - X) ratio of classes

    Returns:
        tps: List of lists, where each "sublist" includes all true positive
            annotations from a document
        fps: List of lists, where each "sublist" includes all false positive
            annotations from a document
        fns: List of lists, where each "sublist" includes all false negative
            annotations from a document
    """
    # Validate inputs
    if tp_ratio and tp_match_type not in {"exact", "overlap"}:
        raise ValueError(
            f"tp_match_type={tp_match_type} not allowed. Must be one of ('exact', 'overlap')"
        )
    if fp_ratio and fp_match_type not in {"exact", "overlap"}:
        raise ValueError(
            f"fp_match_type={fp_match_type} not allowed. Must be one of ('exact', 'overlap')"
        )
    if fn_ratio and fn_match_type not in {"exact", "overlap"}:
        raise ValueError(
            f"fn_match_type={fn_match_type} not allowed. Must be one of ('exact', 'overlap')"
        )

    equality_slc_map = {
        "exact": sequence_labeling_conf_matrix(true, predicted, sequence_exact_match),
        "overlap": sequence_labeling_conf_matrix(true, predicted, sequences_overlap),
    }

    # Get true positives and false positives in the format returned by sequence_labeling_conf_matrix
    slc_tp = equality_slc_map[tp_match_type] if tp_ratio else {}
    slc_fp = equality_slc_map[fp_match_type] if fp_ratio else {}
    slc_fn = equality_slc_map[fn_match_type] if fn_ratio else {}

    label_classes = set(list(slc_fp.keys()) + list(slc_fp.keys()) + list(slc_fn.keys()))
    # Sample labels to keep only a subset of classes if missing_label_ratio is present
    if missing_label_ratio:
        random.seed(random_seed)
        label_classes = random.sample(
            label_classes, int(len(label_classes) * (1 - missing_label_ratio))
        )

    flat_tps, flat_fps, flat_fns, = (
        [],
        [],
        [],
    )
    for label in label_classes:
        if label in slc_tp:
            flat_tps.extend(slc_tp[label]["true_positives"])
        # Sometimes we see predictions with identical start and end indices
        # It does not make sense to label these sequences since they don't
        # span any text, so skip them
        if label in slc_fp:
            flat_fps.extend(
                fp
                for fp in slc_fp[label]["false_positives"]
                if fp["start"] != fp["end"]
            )
        if label in slc_fn:
            flat_fns.extend(slc_fn[label]["false_negatives"])

    # Sample from flattened list(s)
    # Don't need to sample if ratio is 1
    # If ratio is 0, set flat_XXs back to an empty list
    if not tp_ratio:
        flat_tps = []
    elif tp_ratio != 1:
        num_samples = int(tp_ratio * len(flat_tps))
        random.seed(random_seed)
        flat_tps = random.sample(flat_tps, num_samples)

    if not fp_ratio:
        flat_fps = []
    elif fp_ratio != 1:
        num_samples = int(fp_ratio * len(flat_fps))
        random.seed(random_seed)
        flat_fps = random.sample(flat_fps, num_samples)

    if not fn_ratio:
        flat_fns = []
    elif fn_ratio != 1:
        num_samples = int(fn_ratio * len(flat_fns))
        random.seed(random_seed)
        flat_fns = random.sample(flat_fns, num_samples)

    # tps, fps, and fns should contain len(true) number of lists
    tps = [[] for _ in range(len(true))]
    fps = [[] for _ in range(len(true))]
    fns = [[] for _ in range(len(true))]

    for tp in flat_tps:
        tp["type"] = "tp"
        idx = tp["doc_idx"]
        tps[idx].append(tp)

    for fp in flat_fps:
        fp["orig_label"] = fp["label"]
        if not replace_fps:
            # This is flawed, because treating false positives as <PAD> sequences
            # may cause problems if the sequence actually is label A, but was
            # predicted as label B. In this case, we would be teaching the model the
            # incorrect behavior
            fp["label"] = "<PAD>"
        fp["type"] = "fp"
        idx = fp["doc_idx"]
        fps[idx].append(fp)

    for fn in flat_fns:
        fn["type"] = "fn"
        idx = fn["doc_idx"]
        fns[idx].append(fn)

    return tps, fps, fns

###################
# Complex metrics #
###################

def alphanumeric_lower(label_text: str) -> str:
    """
    Given text for a label, return lowercase alphanumeric characters

    Args:
        label_text

    Returns:
        label_text
    """
    return "".join(char for char in label_text if char.isalnum()).lower()

def find_many_match(
    label: Dict[str, Union[int, str]], preds: List[Dict[str, Union[int, str]]]
) -> Union[Dict[str, Union[int, str]], None]:
    """
    Check if any prediction matches with label based on checking for an
    alphanumeric lowercase string match.

    Each prediction can only match with one label (and vice versa), so
    this requires the prediction match candidate to not already be matched
    to a label.

    Returns None if there is no match

    Args:
        label: Ground truth label
        preds: List of predictions

    Returns:
        pred: Matched prediction
    """
    for pred in preds:
        if (
            alphanumeric_lower(label["text"]) == alphanumeric_lower(pred["text"])
        ) and not pred["matched"]:
            return pred

    return None

def many_conf_matrix(
    labels: List[Dict[str, Union[int, str]]], preds: List[Dict[str, Union[int, str]]]
) -> Tuple[int, int, int]:
    """
    Given ground truth labels and predictions for a given class of a document,
    compute the confusion matrix

    Args:
        labels: Labels for a document
        preds: Predictions for a document

    Returns:
        num_tps: Number of true positives for that class and document
        num_fps: : Number of false positives for that class and document
        num_fns: Number of false negatives for that class and document
    """
    # Initialize match values to False
    for label in labels:
        label["matched"] = False
    for pred in preds:
        pred["matched"] = False

    for label in labels:
        pred_match = find_many_match(label, preds)
        if pred_match:
            label["matched"] = True
            pred_match["matched"] = True

    num_tps = sum(int(label["matched"]) for label in labels)
    num_fns = sum(int(not label["matched"]) for label in labels)
    num_fps = sum(int(not pred["matched"]) for pred in preds)
    return num_tps, num_fns, num_fps

def calc_macro_metric(
    one_metric: Dict[str, float], many_metric: Dict[str, float] = {}
) -> float:
    """
    Given per class metrics for recall, precision, or F1, for both classes
    that appear once or many times per doc, take the mean of the classes to
    return a macro metric.

    Args:
        one_metric: Per class recall, precision, or F1 for labels that should
            appear max once per doc
        many_metric: Same as above, but for labels that appear many times per doc

    Returns:
        macro_metric: Macro average of recall, precision, or F1
    """

    return np.array(list(one_metric.values()) + list(many_metric.values())).mean()

def calc_micro_metric(
    one_metric: Dict[str, float],
    one_counts: Dict[str, Dict[str, int]],
    many_metric: Dict[str, float] = None,
    many_counts: Dict[str, Dict[str, int]] = None,
) -> float:
    """
    Given per class metrics for recall, precision, or F1, for both classes
    that appear once or many times per doc, use the also provided counts
    for each class to compute a weighted average that weights each class
    based on the amount of support it has.

    Args:
        one_metric: Per class recall, precision, or F1 for labels that should
            appear max once per doc
        one_counts: Per class support for one_metric labels for weighting
        many_metric: Same as one_metric, but for labels that appear many times per doc
        many_counts: Per class support for many_metric labels for weighting

    Returns:
        micro_metric: Micro average of recall, precision, or F1
    """
    metric_sum = 0
    metric_counts = 0
    for cls in one_metric:
        cls_count = one_counts[cls]["tp"] + one_counts[cls]["fn"]
        metric_sum += one_metric[cls] * cls_count
        metric_counts += cls_count
    if many_metric:
        for cls in many_metric:
            cls_count = many_counts[cls]["tp"] + many_counts[cls]["fn"]
            metric_sum += many_metric[cls] * cls_count
            metric_counts += cls_count
    return metric_sum / metric_counts

def seq_recall_precision_f1(
    true: List[List[Dict[str, Union[int, str]]]],
    predicted: List[List[Dict[str, Union[int, str]]]],
    span_type: str = "token",
) -> Tuple[
    Dict[str, float], Dict[str, float], Dict[str, float], Dict[str, Dict[str, int]]
]:
    """
    Given true and predicted labels, calculate per class precision,
    recall, F1, and and support

    Args:
        true: Sequence labeling labels for each document
        predicted: Model predictions for labels for each doc
        span_type: How to determine matches between true and prediction.
            Valid options are "token", "overlap", "exact", and "superset"

    Returns:
        recall: Per class recall
        precision: Per class precision
        f1: Per class F1
        support: FP, FN, and TP counts per class
    """
    conf_matrix_fn = get_seq_conf_matrix_fn(span_type)
    class_conf_matrices = conf_matrix_fn(true, predicted)
    recall = {}
    precision = {}
    f1 = {}
    support = {}
    for cls, conf_matrix in class_conf_matrices.items():
        num_fns = len(conf_matrix["false_negatives"])
        num_fps = len(conf_matrix["false_positives"])
        num_tps = len(conf_matrix["true_positives"])
        # Skip classes with no support
        if num_fns + num_fps + num_tps == 0:
            continue
        recall[cls] = calc_recall(num_tps, num_fns)
        precision[cls] = calc_precision(num_tps, num_fps)
        f1[cls] = calc_f1(recall[cls], precision[cls])
        support[cls] = {"fp": num_fps, "fn": num_fns, "tp": num_tps}

    return recall, precision, f1, support

def find_best_label(
    label_cls: str,
    label_candidates: List[Dict[str, Union[int, str]]],
    doc_fns: List[Dict[str, Union[int, str]]],
) -> str:
    """
    Helper function for dedup_test_labels() that is called per label per document
    any time there are multiple labels of a class that should only appear once
    per document.

    Args:
        label_cls: Name of target label class
        label_candidates: Labels of the target label class
        doc_fns: All false negatives for the current document

    Returns:
        text: Text for the chosen best label
    """
    # First try to find false negative
    for fn in doc_fns:
        if fn["label"] == label_cls:
            return fn["text"]
    # If no false negatives, default to first label
    sorted_cands = sorted(label_candidates, key=lambda x: x["start"])
    return sorted_cands[0]["text"]

def dedup_test_labels(
    doc_labels: Dict[str, List[Dict[str, Union[int, str]]]],
    doc_fns: List[Dict[str, Union[int, str]]],
) -> Dict[str, str]:
    """
    For classes that should only appear once per document, we need a way
    to choose the "best" label from review data in the case of duplicates.
    The methodology for doing so is as follows.

    1. If there are any false negatives for this class, choose the first
        one (by start index)
    2. Else, use the first value in the review set (by start index)

    Args:
        doc_labels: Dictionary with a list per label class with
            candidates for that label
        doc_fns: List of all false negatives for a doc

    Returns:
        deduped_labels: Dictionary with the best label for each class
    """
    deduped_labels = {}
    doc_fns = sorted(doc_fns, key=lambda x: x["start"])
    for label_cls, labels in doc_labels.items():
        if len(labels) == 1:
            deduped_labels[label_cls] = labels[0]["text"]
        else:
            deduped_labels[label_cls] = find_best_label(label_cls, labels, doc_fns)
    return deduped_labels

def run_evaluation(
    test_preds: List[List[Dict[str, Union[int, str]]]],
    test_labels: List[List[Dict[str, Union[int, str]]]],
    one_classes: Set[str],
    verbose: bool = False,
) -> Dict[str, float]:
    """
    Complex evaluation methodology to treat labels separately depending on if they
    should appear at most once per document or several times. This methodology
    assumes that we are using review data for our test set.
    \
    1. Separate test preds into one vs many
    2. For one, process down to a single value (no indices) based on the
        highest confidence prediction
    3. Separate test labels into one vs many
    4. Get FNs b/w test preds and labels
    5. For one, process down to max one label per document in this order:
        a) If there are any false negatives, choose the first one (by start index)
        b) Use the first value in the labels (by start index)
    6. Use existing metrics calculations to get per class precision and recall
        on "many" preds vs labels
    7. Use string matching (only alphanumeric chars, all lowercase) to determine
        TPs, FPs, and FNs for "one" preds vs labels
    8. Calculate precision and recall on "one" classes
    9. Calculate macro metrics
    10. Calculate micro metrics

    Args:
        test_preds: Predictions on the test set for each document
        test_labels: Labels on the test set for each document
        one_classes: Set of class names that should appear max once per doc
        verbose: If True, print more detailed recall/precision/F1 statistics
    """
    # 1. Separate test preds into one vs many
    # 2. For one, process down to a single value (no indices) based on the
    # highest confidence prediction
    # Note that test_preds_one is a list of dictionaries (one entry per label
    # per doc) while test_preds_many is a list of lists
    test_preds_one = []
    test_preds_many = []
    for doc_preds in test_preds:
        doc_preds_one = {}
        doc_preds_many = []
        for pred in doc_preds:
            label = pred["label"]
            if label in one_classes:
                if label in doc_preds_one:
                    # Replace existing pred in doc_preds_one if the confidence
                    # of the current pred is higher
                    conf = pred["confidence"][label]
                    if conf > doc_preds_one[label]["conf"]:
                        doc_preds_one[label] = {"text": pred["text"], "conf": conf}
                else:
                    doc_preds_one[label] = {
                        "text": pred["text"],
                        "conf": pred["confidence"][label],
                    }
            else:
                doc_preds_many.append(pred)
        test_preds_one.append(doc_preds_one)
        test_preds_many.append(doc_preds_many)

    # 3. Separate test labels into one vs many
    test_labels_one_w_dups = []
    test_labels_many = []
    for doc_labels in test_labels:
        doc_label_candidates_one = {}
        doc_labels_many = []
        for label in doc_labels:
            if label["label"] in one_classes:
                if label["label"] in doc_label_candidates_one:
                    doc_label_candidates_one[label["label"]].append(label)
                else:
                    doc_label_candidates_one[label["label"]] = [label]
            else:
                doc_labels_many.append(label)
        test_labels_one_w_dups.append(doc_label_candidates_one)
        test_labels_many.append(doc_labels_many)

    # 4. Get FNs between test preds and labels
    _, _, test_fns = identify_tps_fps_fns(
        true=test_labels,
        predicted=test_preds,
        tp_ratio=0,
        fp_ratio=0,
        fn_ratio=1,
    )

    # 5. For "one", process down to max one label per document
    test_labels_one = []
    for doc_labels, doc_fns in zip(test_labels_one_w_dups, test_fns):
        test_labels_one.append(dedup_test_labels(doc_labels, doc_fns))

    # 6. Use existing metrics calculations to get per class precision and recall
    # on "many" preds vs labels
    many_recall, many_precision, many_f1, many_counts = seq_recall_precision_f1(
        test_labels_many, test_preds_many
    )

    # 7. Use string matching (only alphanumeric chars, all lowercase) to determine
    # TPs, FPs, and FNs for "one" preds vs labels
    cm_counts_one = {cls: defaultdict(int) for cls in one_classes}
    for doc_preds, doc_labels in zip(test_preds_one, test_labels_one):
        for cls in one_classes:
            pred_text = doc_preds.get(cls, {}).get("text")
            label_text = doc_labels.get(cls, {}).get("text")
            if pred_text is not None and label_text is None:
                cm_counts_one[cls]["fp"] += 1
            elif label_text is not None and pred_text is None:
                cm_counts_one[cls]["fn"] += 1
            elif pred_text is not None and label_text is not None:
                # Process pred and label to lowercase alphanumeric chars for comparison
                pred = alphanumeric_lower(pred_text)
                label = alphanumeric_lower(label_text)
                if pred == label:
                    cm_counts_one[cls]["tp"] += 1
                else:
                    cm_counts_one[cls]["fp"] += 1
                    cm_counts_one[cls]["fn"] += 1

    # 8. Calculate precision and recall on "one" classes
    one_recall = {}
    one_precision = {}
    one_f1 = {}
    for cls, counts in cm_counts_one.items():
        # Skip classes with no support if that happens for any reason
        if counts["tp"] + counts["fp"] + counts["fn"] > 0:
            one_recall[cls] = calc_recall(counts["tp"], counts["fn"])
            one_precision[cls] = calc_precision(counts["tp"], counts["fp"])
            one_f1[cls] = calc_f1(one_recall[cls], one_precision[cls])

    # 9. Macro metrics
    macro_recall = calc_macro_metric(one_recall, many_recall)
    macro_precision = calc_macro_metric(one_precision, many_precision)
    macro_f1 = calc_macro_metric(one_f1, many_f1)

    # 10. Micro metrics
    micro_recall = calc_micro_metric(
        one_recall, cm_counts_one, many_recall, many_counts
    )
    micro_precision = calc_micro_metric(
        one_precision, cm_counts_one, many_precision, many_counts
    )
    micro_f1 = calc_micro_metric(one_f1, cm_counts_one, many_f1, many_counts)

    if verbose:
        print("One recall", one_recall)
        print("One precision", one_precision)
        print("One f1", one_f1)
        print("One support", cm_counts_one)
        print("Many recall", many_recall)
        print("Many precision", many_precision)
        print("Many f1", many_f1)
        print("Many support", many_counts)

    final_metrics = {
        "macro_recall": macro_recall,
        "macro_precision": macro_precision,
        "macro_f1": macro_f1,
        "micro_recall": micro_recall,
        "micro_precision": micro_precision,
        "micro_f1": micro_f1,
    }
    return final_metrics

def eval_value_matching(
    test_preds: List[List[Dict[str, Union[int, str]]]],
    test_labels: List[List[Dict[str, Union[int, str]]]],
    one_classes: Set[str],
    all_classes: Set[str],
    verbose: bool = False,
) -> Dict[str, float]:
    """
    Complex evaluation methodology to treat labels separately depending on if they
    should appear at most once per document or several times.

    This function is identical to run_evaluation(), except that for this comparison,
    we compare values directly instead of using start/end indices.

    1. Separate test preds into one vs many
    2. For one, process down to a single value (no indices) based on the
        highest confidence prediction
    3. Separate test labels into one vs many
    4. Get FNs between test preds and labels by doing string matches
    5. For "one", process each class down to a set of labels using the following logic:
        a) If there are any false negatives, create a set of them
        b) Else, create a set of all of the labels of that class
    6. For "many" preds vs labels, use string matching (only alphanumeric chars,
        all lowercase) to get TPs, FPs, and FNs
    7. Calculate precision and recall on "many" classes
    8. Apply the same string matching technique to "one" preds vs labels
    9. Calculate precision and recall on "one" classes
    10. Calculate micro metrics
    11. Calculate macro metrics

    Args:
        test_preds: Predictions on the test set for each document
        test_labels: Labels on the test set for each document
        one_classes: Set of class names that should appear max once per doc
        all_classes: Set of all classes that can appear for this workflow
        verbose: If True, print more detailed recall/precision/F1 statistics

    Returns:
        eval_metrics: Micro and macro precision, recall, and F1
    """
    # 1. Separate test preds into one vs many
    # 2. For one, process down to a single value (no indices) for each class
    # based on the highest confidence prediction
    # For many, process down to a list of labels for each class
    test_preds_one = []
    test_preds_many = []
    for doc_preds in test_preds:
        doc_preds_one = {}
        doc_preds_many = defaultdict(list)
        for pred in doc_preds:
            label = pred["label"]
            if label in one_classes:
                if label in doc_preds_one:
                    # Replace existing pred in doc_preds_one if the confidence
                    # of the current pred is higher
                    conf = pred["confidence"][label]
                    if conf > doc_preds_one[label]["conf"]:
                        doc_preds_one[label] = {"text": pred["text"], "conf": conf}
                else:
                    doc_preds_one[label] = {
                        "text": pred["text"],
                        "conf": pred["confidence"][label],
                    }
            else:
                doc_preds_many[label].append(pred)
        test_preds_one.append(doc_preds_one)
        test_preds_many.append(doc_preds_many)

    # 3. Separate test labels into one vs many
    # "one" has a set for each class for each doc. Each element of the set
    #   is just alphanumeric lowercase text
    # "many" has a list of the full label dicts for each class in each doc
    test_labels_one_w_dups = []
    test_labels_many = []
    for doc_labels in test_labels:
        doc_label_candidates_one = defaultdict(set)
        doc_labels_many = defaultdict(list)
        for label in doc_labels:
            if label["label"] in one_classes:
                doc_label_candidates_one[label["label"]].add(
                    alphanumeric_lower(label["text"])
                )
            else:
                doc_labels_many[label["label"]].append(label)
        test_labels_one_w_dups.append(doc_label_candidates_one)
        test_labels_many.append(doc_labels_many)

    # 4. Get FNs between test preds and labels by doing string matches
    # Each element of test_fns is a defaultdict from label class name
    # to a set of the false negatives for the label in the current doc
    test_fns = []
    for doc_labels, doc_preds in zip(test_labels, test_preds):
        doc_fns = defaultdict(set)
        doc_pred_strs = set(alphanumeric_lower(pred["text"]) for pred in doc_preds)
        for label in doc_labels:
            label_str = alphanumeric_lower(label["text"])
            # If this condition is met, label is a false negative
            if label_str not in doc_pred_strs:
                doc_fns[label["label"]].add(label_str)
        test_fns.append(doc_fns)

    # 5. For "one", process each class down to a set of labels using the following logic:
    #   a) If there are any false negatives, create a set of them
    #   b) Else, create a set of all of the labels of that class
    test_labels_one = []
    for doc_labels, doc_fns in zip(test_labels_one_w_dups, test_fns):
        doc_labels_one = {}
        for cls in one_classes:
            if doc_fns.get(cls):
                doc_labels_one[cls] = doc_fns[cls]
            elif doc_labels_one.get(cls):
                doc_labels_one[cls] = doc_labels[cls]
        test_labels_one.append(doc_labels_one)

    # 6. For "many" preds vs labels, use string matching (only alphanumeric chars,
    # all lowercase) to get TPs, FPs, and FNs
    many_classes = all_classes - one_classes
    cm_counts_many = {cls: defaultdict(int) for cls in many_classes}
    for doc_preds, doc_labels in zip(test_preds_many, test_labels_many):
        for cls in many_classes:
            tps, fps, fns = many_conf_matrix(
                doc_labels.get(cls, []), doc_preds.get(cls, [])
            )
            cm_counts_many[cls]["tp"] += tps
            cm_counts_many[cls]["fp"] += fps
            cm_counts_many[cls]["fn"] += fns

    # 7. Calculate precision and recall on "many" classes
    many_recall = {}
    many_precision = {}
    many_f1 = {}
    for cls, counts in cm_counts_many.items():
        # Skip classes with no support if that happens for any reason
        if counts["tp"] + counts["fp"] + counts["fn"] > 0:
            many_recall[cls] = calc_recall(counts["tp"], counts["fn"])
            many_precision[cls] = calc_precision(counts["tp"], counts["fp"])
            many_f1[cls] = calc_f1(many_recall[cls], many_precision[cls])

    # 8. Use string matching (only alphanumeric chars, all lowercase) to determine
    # TPs, FPs, and FNs for "one" preds vs labels
    cm_counts_one = {cls: defaultdict(int) for cls in one_classes}
    for doc_preds, doc_labels in zip(test_preds_one, test_labels_one):
        for cls in one_classes:
            pred = doc_preds.get(cls)
            labels = doc_labels.get(cls, {})
            if pred is not None and labels == {}:
                cm_counts_one[cls]["fp"] += 1
            elif pred is None and labels != {}:
                cm_counts_one[cls]["fn"] += 1
            elif pred is not None and labels != {}:
                # Check if pred is in the set of labels or not
                if alphanumeric_lower(pred["text"]) in labels:
                    cm_counts_one[cls]["tp"] += 1
                # Even if the set of labels has multiple "wrong" labels, we only
                # treat it as a single error, because ideally we would have a
                # constraint in place to only have a single prediction for labels
                # that should appear max once per document
                else:
                    cm_counts_one[cls]["fp"] += 1
                    cm_counts_one[cls]["fn"] += 1

    # 9. Calculate precision and recall on "one" classes
    one_recall = {}
    one_precision = {}
    one_f1 = {}
    for cls, counts in cm_counts_one.items():
        # Skip classes with no support if that happens for any reason
        if counts["tp"] + counts["fp"] + counts["fn"] > 0:
            one_recall[cls] = calc_recall(counts["tp"], counts["fn"])
            one_precision[cls] = calc_precision(counts["tp"], counts["fp"])
            one_f1[cls] = calc_f1(one_recall[cls], one_precision[cls])

    # 10. Macro metrics
    macro_recall = calc_macro_metric(one_recall, many_recall)
    macro_precision = calc_macro_metric(one_precision, many_precision)
    macro_f1 = calc_macro_metric(one_f1, many_f1)

    # 11. Micro metrics
    micro_recall = calc_micro_metric(
        one_recall, cm_counts_one, many_recall, cm_counts_many
    )
    micro_precision = calc_micro_metric(
        one_precision, cm_counts_one, many_precision, cm_counts_many
    )
    micro_f1 = calc_micro_metric(one_f1, cm_counts_one, many_f1, cm_counts_many)

    if verbose:
        print("One recall", one_recall)
        print("One precision", one_precision)
        print("One f1", one_f1)
        print("One support", cm_counts_one)
        print("Many recall", many_recall)
        print("Many precision", many_precision)
        print("Many f1", many_f1)
        print("Many support", cm_counts_many)

    final_metrics = {
        "macro_recall": macro_recall,
        "macro_precision": macro_precision,
        "macro_f1": macro_f1,
        "micro_recall": micro_recall,
        "micro_precision": micro_precision,
        "micro_f1": micro_f1,
    }
    return final_metrics


from PyQt5 import QtCore, QtGui
from PyQt5.QtWidgets import QApplication, QWidget
import sys

class Example(QWidget):
    def __init__(self):
        super(Example, self).__init__()
        self.setGeometry(300, 300, 250, 150)
        self.show()

    def keyPressEvent(self, event):
        if event.key() == QtCore.Qt.Key_Q:
            print ("Killing")
            self.deleteLater()
        elif event.key() == QtCore.Qt.Key_Enter:
            self.proceed()
        event.accept()

    def proceed(self):
        print ("Call Enter Key")

def main():
    app = QApplication(sys.argv)
    ex = Example()
    sys.exit(app.exec_())

if __name__ == '__main__':
    main()


"""
VUnit testbench for the High Speed ADC AXIS Adapter
"""

from pathlib import Path

import vunit_util

WORKSPACE = Path(__file__).parent / "workspace"
RTL_ROOT = Path(__file__).parent / ".."

# Create source list
sources = [
    RTL_ROOT / "testbenches" / "hsadc_axis_wrapper_tb.sv",
    RTL_ROOT / "hsadc" / "*.sv",
    RTL_ROOT / "axis" / "*.sv"
]

# Create VUnit instance and add sources to library
vu, lib = vunit_util.init(WORKSPACE)
lib.add_source_files(sources)

# Create testbench
tb = lib.test_bench("hsadc_axis_wrapper_tb")

# Suppress vopt deprecation error
vu.add_compile_option('modelsim.vlog_flags', ['-suppress', '12110'])

# Run
vu.main()


# 1 - https://www.qwant.com/?client=brz-moz&t=videos&q=pygame&o=0%3AWxEmflz009U
# 2 - https://www.youtube.com/watch?v=W4cbtcyONag

import sys
import pygame

pygame.init()

# Creating colors
BLUE = (0, 0, 255)
RED = (255, 0, 0)
GREEN = (0, 255, 0)
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)

screen = pygame.display.set_mode((400, 400))
screen.fill(WHITE)

image = pygame.image.load("ball.png")
image = pygame.transform.rotozoom(image, 0, 0.2)

x = 0
y = 0

clock = pygame.time.Clock()

while True:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            pygame.quit()
            sys.exit()

    pressed = pygame.key.get_pressed()
    if pressed[pygame.K_LEFT]:
        x -= 1
    if pressed[pygame.K_RIGHT]:
        x += 1
    if pressed[pygame.K_UP]:
        y -= 1
    if pressed[pygame.K_DOWN]:
        y += 1

    """
    On repeint en blanc le fond afin d'effacer les traces laisser par le deplacement du ballon.
    """
    screen.fill(WHITE)
    screen.blit(image, (x, y))
    pygame.display.flip()
    clock.tick(60)

"""
pygame.display.update()
Les deux fonctions font la meme action elles mettent à jour les informations de l'écran en appliquant les
changements faits. Cependant pygame.display.update() a une particularité, à l'aide d'un tuple de 4 valeurs
(x, y. width, height) ou d'un objet rect: pygame.Rect(x, y, width, height) nous pouvons définir la zone que nous
souhaitons mettre à jour. Ainsi si nous utilisons la méthode sous cette forme pygame.display.update((0, 0, 150, 150))
nous mettons à jour qu'un carré de 150 pixels par 150 pixels qui se trouve aux coordonnées 0, 0.
pygame.display.flip() quand à lui ne prend aucun paramètre et permet de mettre à jour l'écran entier. Cependant
flip a été codé avec certaine optimisations qui permettent de mettre à jour l'écran plus rapidement.
"""

"""
pygame en oriente objet
"""


import torch
import numpy as np

'''
Train the model on the given train set for one epoch
Args
model:       the model we want to train
train_loader:a dataloader object for the train set
optimizer:   the optimizer to use
loss_fn:     the loss function to use
device:      the device to use (CPU or cuda)
print_every: a boolean telling when to print infos
status:      a boolean indicating debug mode
Return
np.mean(np.array(losses)):  the average loss
accuracy:                   the average accuracy
'''
def train(model, train_loader, optimizer, loss_fn, device, print_every=10, status = False):
    model.train()
    losses = []
    n_correct = 0
    n_sample = 0
    n_sample_0 = 0
    n_sample_1 = 0
    iteration = 0
    print("length of dataset :",len(train_loader.dataset))
    for iter, (images, labels) in enumerate(train_loader):
        images = images.to(device)
        labels = labels.to(device)
        output = model(images)
        optimizer.zero_grad()
        loss = loss_fn(output, labels)
        loss.backward()
        optimizer.step()

        losses.append(loss.item())
        n_correct += torch.sum(output.argmax(1) == labels).item()

        # DEBUG
        if status is True:
            n_sample_0 += torch.sum(output.argmax(1) == 0).item()
            n_sample_1 += torch.sum(output.argmax(1) == 1).item()
            n_sample = n_sample_0+n_sample_1
            if(n_sample>0):
                accuracy = 100.0 * n_correct / n_sample
            else:
                accuracy = np.inf

            print('Training iteration {}: loss {:.4f}, accuracy {:.4f}'.format(iteration, loss.item(), accuracy))
            print('# 0 output: {}, # 1 output: {}'.format(n_sample_0,n_sample_1))

        iteration += 1

    accuracy = 100.0 * n_correct / len(train_loader.dataset)

    return np.mean(np.array(losses)), accuracy


'''
Created on 4 de dez. de 2020

@author: michel
'''

import csv
from datetime import datetime
import get_date

## Getting the Dates
get_date.get_dp_date()

with open('./input_data/included-datapoints-date.csv', newline='') as csvfile:
    with open('./output_data/energy-datapoints.csv',  mode='w') as csvfileout:
        writer = csv.writer(csvfileout, delimiter=',', quotechar='"', quoting=csv.QUOTE_MINIMAL)
        reader = csv.DictReader(csvfile)
        # Headers
        line = []
        line.append('ID')
        line.append('URL')
        line.append('Collection')
        line.append('Title')
        line.append('Contents')
        line.append('DateTime')
        line.append('Timestamp')
        writer.writerow(line)
        for row in reader:
            line = []
            line.append(row['ID'])
            line.append(row['URL'])
            line.append(row['Collection'])
            line.append(row['Title'])
            line.append(row['Contents'])
            line.append(row['DateTime'])
            date = row['DateTime']
            dtime = '0000-00-00 00:00:00'
            timestamp = '0'
            collection = row['Collection']
            if collection == 'Commits':
                if date != '0000-00-00 00:00:00' and date != '0000-00-00 00:00':
                    dtime = datetime.strptime(date, '%Y-%m-%d %H:%M:%S')
                    #print('Commits', dtime)
                    timestamp = datetime.timestamp(dtime)
            elif collection in ['GitHubIssues', 'GitHubPRs', 'Repositories']:
                if date != 'NA':
                    dtime = datetime.strptime(date, '%b %d, %Y')
                    #print(collection, dtime)
                    timestamp = datetime.timestamp(dtime)
            elif collection in ['ROSAnswers','ROSDiscourse','StackOverflow']:
                date_parts = date.split()
                if collection == 'ROSAnswers':
                    date = str(date_parts[0])+' '+str(date_parts[1])
                else:
                    date = str(date_parts[0])+' '+str(date_parts[1])+':00'
                #print(collection,date)
                dtime = datetime.strptime(date, '%Y-%m-%d %H:%M:%S')
                timestamp = datetime.timestamp(dtime)
            line.append(timestamp)
            print(timestamp)
            writer.writerow(line)


#!/usr/bin/env python
#coding: utf-8
from __future__ import print_function

# PART I

# PART 2

# PART 3

# PART 4
'''
from selenium import webdriver

class ClassName(object):
	"""docstring for ClassName"""
	def __init__(self, arg):
		super(ClassName, self).__init__()
		self.arg = arg

'''

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

#!/usr/bin/env python
#coding: utf-8
#from __future__ import print_function

'''
Class 2:

# 导入 selenium
from selenium import webdriver

# 打开谷歌浏览器，并访问百度
browser = webdriver.Chrome()
browser.get('http://www.baidu.com/')
'''

'''
Class 2:

from selenium import webdriver
from selenium.webdriver.common.keys import Keys

driver = webdriver.Chrome()
driver.get("http://www.baidu.com")

# 获取浏览器中 name 为 wd 的标签
elem = driver.find_element_by_name("wd")
# 搜索 cxy61
elem.send_keys("cxy61")
elem.send_keys(Keys.RETURN)
# 打印页面
print driver.page_source
##driver.quit()
'''

# PART  III,  Class 3
'''

from selenium import webdriver

from selenium.webdriver.common.keys import Keys

dr = webdriver.Chrome()
dr.get('https://www.qiushibaike.com/')

# 获取 id 为 content-left 的元素
main_content = dr.find_element_by_id('content-left')
# 获取 class 为 content 的元素
contents = main_content.find_elements_by_class_name('content')

# for 循环
i = 1
for content in contents:
	#print(content.text)
	print(str(i)+'.' + content.text + '\n')
	i += 1

# 退出
dr.quit()

'''

from selenium import webdriver

class Qiubai:
	def __init__(self):
		self.dr = webdriver.Chrome()
		self.dr.get('https://www.qiushibaike.com/')

	def print_contenx(self):
		main_content = self.dr.find_element_by_id('content-left')
		contents = main_content.find_elements_by_class_name('content')

		i = 1
		for content in contents:
			print(str(i)+"." + content.text + '\n')
			i += 1

		self.quit()

	def quit(self):
		self.dr.quit()

Qiubai().print_contenx()


from __future__ import absolute_import
from __future__ import division
# from __future__ import google_type_annotations
from __future__ import print_function
from albert import fine_tuning_utils
from albert import modeling
from albert import optimization
from albert import tokenization
import tensorflow.compat.v1 as tf
from tensorflow.contrib import data as contrib_data
import json
import collections

class InputExample(object):
  """A single training/test example for crf based models."""

  def __init__(self, guid, text, token_labels=None):
    """Constructs a InputExample.

    Args:
      guid: Unique id for the example.
      text_a: string. The untokenized text of the first sequence.
      token_labels: (Optional) string. for each token in text.
    """
    self.guid = guid
    self.text = text
    self.token_labels=token_labels
    #if self.token_labels is None:
    #    self.token_labels = ["O"]* 20

class PaddingInputExample(object):
  """Fake example so the num input examples is a multiple of the batch size.

  When running eval/predict on the TPU, we need to pad the number of examples
  to be a multiple of the batch size, because the TPU requires a fixed batch
  size. The alternative is to drop the last batch, which is bad because it means
  the entire output data won't be generated.

  We use this class instead of `None` because treating `None` as padding
  battches could cause silent errors.
  """

class InputFeatures(object):
  """A single set of features of data."""

  def __init__(self,
               input_ids,
               input_mask,
               segment_ids,
               label_ids,
               guid=None,
               example_id=None,
               is_real_example=True):
    self.input_ids = input_ids
    self.input_mask = input_mask
    self.segment_ids = segment_ids
    self.label_ids = label_ids
    self.example_id = example_id
    self.guid = guid
    self.is_real_example = is_real_example

def load_labels(label_list_file):
    label_list=[]
    with open(label_list_file, "r", encoding="utf-8") as f:
        for line in f:
            if line and line.strip():
                label_list.append(line.strip())
    return label_list

def convert_single_example(ex_index, example, max_seq_length,
                           processor):
    """Converts a single `InputExample` into a single `InputFeatures`."""
    tokenizer=processor.tokenizer
    label_map=processor.label_map

    if isinstance(example, PaddingInputExample):
        return InputFeatures(
            input_ids=[0] * max_seq_length,
            input_mask=[0] * max_seq_length,
            segment_ids=[0] * max_seq_length,
            label_ids=[0] * max_seq_length,
            is_real_example=False)

    # The convention in ALBERT is:
    # (a) For sequence pairs:
    #  tokens:   [CLS] is this jack ##son ##ville ? [SEP] no it is not . [SEP]
    #  type_ids: 0     0  0    0    0     0       0 0     1  1  1  1   1 1
    # (b) For single sequences:
    #  tokens:   [CLS] the dog is hairy . [SEP]
    #  type_ids: 0     0   0   0  0     0 0
    #
    # Where "type_ids" are used to indicate whether this is the first
    # sequence or the second sequence. The embedding vectors for `type=0` and
    # `type=1` were learned during pre-training and are added to the
    # embedding vector (and position vector). This is not *strictly* necessary
    # since the [SEP] token unambiguously separates the sequences, but it makes
    # it easier for the model to learn the concept of sequences.
    #
    # For classification tasks, the first vector (corresponding to [CLS]) is
    # used as the "sentence vector". Note that this only makes sense because
    # the entire model is fine-tuned.

    #tokens = tokenizer.tokenize(example.text)
    tokens=[]
    segment_ids = []
    label_ids=[]

    '''
    a severe bug exits using the following snippets as the original text are labeled
    character by charater ...
    we need to modify it
    it=iter(example.token_labels)

    for token in tokens[1:]:
        segment_ids.append(0)
        if not token.startswith(piece_prefix): # a new token rather than pieces
            label_id=label_map[next(it)]
        label_ids.append(label_id)
    '''

    '''
    a severe error also occured as pieces like cm, etc. take 2 positions chars
    offset=0
    for token in tokenizer.basic_tokenizer.tokenize(example.text):
        if example.token_labels is None:
            label_str="O"
        else:
            label_str=example.token_labels[offset]

        for sub_token in  tokenizer.wordpiece_tokenizer.tokenize(token):
            tokens.append(sub_token)
            label_ids.append(label_map[label_str])
            segment_ids.append(0)

        offset+=len(token.replace(piece_prefix,""))
    '''

    '''
    tokstr="".join(map(lambda tok: tok.replace(piece_prefix, ""),tokens[1:]))
    print("-------", tokstr)
    print("_______", example.text)
    print("******", example.token_labels)
    print(len(tokstr),len(example.token_labels),len(example.text))
    '''

    text=example.text
    for offset, word in enumerate(text):
        if example.token_labels is None:
            label_str="O"
        else:
            label_str=example.token_labels[offset]

        label=label_map[label_str]
        if word not in tokenizer.vocab:
            word="[UNK]"

        tokens.append(word )
        segment_ids.append(0)
        label_ids.append(label)

    # Account for [CLS] and [SEP] with "- 2"
    if len(tokens) > max_seq_length - 2:
        tokens = tokens[0:(max_seq_length - 2)]
        segment_ids=segment_ids[0:(max_seq_length-2)]
        label_ids=label_ids[0:(max_seq_length-2)]

    # inser [CLS]
    tokens.insert(0,"[CLS]")
    segment_ids.insert(0,0)
    label_ids.insert(0,label_map["O"])
    #add [SEP]
    tokens.append("[SEP]")
    segment_ids.append(0)
    label_ids.append(label_map["O"])

    input_ids = tokenizer.convert_tokens_to_ids(tokens)

    # The mask has 1 for real tokens and 0 for padding tokens. Only real
    # tokens are attended to.
    input_mask = [1] * len(input_ids)

    #print(len(input_ids), len(input_mask), len(segment_ids), len(label_ids))
    # Zero-pad up to the sequence length.
    while len(input_ids) < max_seq_length:
        input_ids.append(0)
        input_mask.append(0)
        segment_ids.append(0)
        label_ids.append(label_map["O"])
    #print(len(input_ids), len(input_mask), len(segment_ids), len(label_ids))
    assert len(input_ids) == max_seq_length
    assert len(input_mask) == max_seq_length
    assert len(segment_ids) == max_seq_length
    assert len(label_ids) == max_seq_length

    if ex_index < 5:
        tf.logging.info("*** Example ***")
        tf.logging.info("guid: %s" % (example.guid))
        tf.logging.info("tokens: %s" % " ".join(
            [tokenization.printable_text(x) for x in tokens]))
        tf.logging.info("input_ids: %s" % " ".join([str(x) for x in input_ids]))
        tf.logging.info("input_mask: %s" % " ".join([str(x) for x in input_mask]))
        tf.logging.info("segment_ids: %s" % " ".join([str(x) for x in segment_ids]))
        if example.token_labels is not None:
            tf.logging.info("label_ids: %s \n   (ids = %s)" % (" ".join(example.token_labels),
            " ".join([str(label_id) for label_id in label_ids]) ) )

    feature = InputFeatures(
        input_ids=input_ids,
        input_mask=input_mask,
        segment_ids=segment_ids,
        label_ids=label_ids,
        guid=example.guid,
        example_id=ex_index,
        is_real_example=True)
    return feature

def data_based_input_fn_builder(features, seq_length, is_training, drop_remainder):
    """Creates an `input_fn` closure to be passed to TPUEstimator."""

    all_input_ids = []
    all_input_mask = []
    all_segment_ids = []
    all_label_ids = []

    for feature in features:
        all_input_ids.append(feature.input_ids)
        all_input_mask.append(feature.input_mask)
        all_segment_ids.append(feature.segment_ids)
        all_label_ids.append(feature.label_ids)

    def input_fn(params):
        """The actual input function."""
        batch_size = params["batch_size"]

        num_examples = len(features)

        # This is for demo purposes and does NOT scale to large data sets. We do
        # not use Dataset.from_generator() because that uses tf.py_func which is
        # not TPU compatible. The right way to load data is with TFRecordReader.
        d = tf.data.Dataset.from_tensor_slices({
            "input_ids":
                tf.constant(
                    all_input_ids, shape=[num_examples, seq_length],
                    dtype=tf.int32),
            "input_mask":
                tf.constant(
                    all_input_mask,
                    shape=[num_examples, seq_length],
                    dtype=tf.int32),
            "segment_ids":
                tf.constant(
                    all_segment_ids,
                    shape=[num_examples, seq_length],
                    dtype=tf.int32),
            "label_ids":
                tf.constant(all_label_ids, shape=[num_examples, seq_length], dtype=tf.int32),
        })

        if is_training:
            d = d.repeat()
            d = d.shuffle(buffer_size=100)

        d = d.batch(batch_size=batch_size, drop_remainder=drop_remainder)
        return d

    return input_fn

def file_based_convert_examples_to_features(
    examples, max_seq_length, seq_processor, output_file):
  """Convert a set of `InputExample`s to a TFRecord file."""

  writer = tf.python_io.TFRecordWriter(output_file)

  for (ex_index, example) in enumerate(examples):
    if ex_index % 10000 == 0:
      tf.logging.info("Writing example %d of %d" % (ex_index, len(examples)))
    #print("-----examples-------\n{}".format(example.token_labels))
    feature = convert_single_example(ex_index, example,
                                     max_seq_length, seq_processor)

    def create_int_feature(values):
      f = tf.train.Feature(int64_list=tf.train.Int64List(value=list(values)))
      return f

    def create_float_feature(values):
      f = tf.train.Feature(float_list=tf.train.FloatList(value=list(values)))
      return f

    features = collections.OrderedDict()
    features["input_ids"] = create_int_feature(feature.input_ids)
    features["input_mask"] = create_int_feature(feature.input_mask)
    features["segment_ids"] = create_int_feature(feature.segment_ids)
    features["label_ids"] =  create_int_feature(feature.label_ids)
    features["is_real_example"] = create_int_feature(
        [int(feature.is_real_example)])

    tf_example = tf.train.Example(features=tf.train.Features(feature=features))
    writer.write(tf_example.SerializeToString())
  writer.close()

def file_based_input_fn_builder(input_file, seq_length, is_training,
                                drop_remainder, bsz,
                                multiple=1):
  """Creates an `input_fn` closure to be passed to TPUEstimator."""
  labeltype = tf.int64

  name_to_features = {
      "input_ids": tf.FixedLenFeature([seq_length * multiple], tf.int64),
      "input_mask": tf.FixedLenFeature([seq_length * multiple], tf.int64),
      "segment_ids": tf.FixedLenFeature([seq_length * multiple], tf.int64),
      "label_ids": tf.FixedLenFeature([seq_length * multiple], labeltype),
      "is_real_example": tf.FixedLenFeature([], tf.int64),
  }

  def _decode_record(record, name_to_features):
    """Decodes a record to a TensorFlow example."""
    example = tf.parse_single_example(record, name_to_features)

    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name in list(example.keys()):
      t = example[name]
      if t.dtype == tf.int64:
        t = tf.to_int32(t)
      example[name] = t

    return example

  def input_fn(params):
    """The actual input function."""
    batch_size = bsz

    # For training, we want a lot of parallel reading and shuffling.
    # For eval, we want no shuffling and parallel reading doesn't matter.
    d = tf.data.TFRecordDataset(input_file)
    if is_training:
      d = d.repeat()
      d = d.shuffle(buffer_size=100)

    d = d.apply(
        contrib_data.map_and_batch(
            lambda record: _decode_record(record, name_to_features),
            batch_size=batch_size,
            drop_remainder=drop_remainder))

    return d

  return input_fn


import codecs
import hashlib
import logging
import re
import time
from datetime import datetime
from typing import List, Union

import requests

from .constants import (CLEVERBOT_API_URL, CLEVERBOT_COOKIE_URL, CLEVERBOT_URL,
                        DEBOUNCE_TIME, MAX_ATTEPTS, MAX_CONTEXT_LENGTH,
                        MAX_DEBOUNCE_ATTEMPS)

class Cleverbot:
    """
    Cleverbot api wrapper

    Attributes
    ----------
    cookies : dict
        Cookies to be used in the requests.
    context : list
        List of messages representing the context of the conversation.
    proxies :
        Proxy to be used or proxies to rotate through.
    debounce: bool
        Whether to try to debounce the requests or not. If set to False it will try to use proxies immediately upon failure.
    :param use_tor_fallback:
        If set to True, whenever a request fail will try to use a new tor session as a proxy. This will override the proxy options. Even when this is set to true it will still try to make a connection without using tor at first.
    Methods
    ----------
    send(message)
        Sends a message to Cleverbot and returns the response as a string.
    """

    def __init__(
        self,
        context: List[str] = [],
        proxies: Union[dict, List[dict], List[str]] = None,
        debounce: bool = True,
        use_tor_fallback: bool = False,
    ):
        """
        Cleverbot Constructor

        :param context List[str]:List of strings to be used as the initial chat messages with the bot.
        :param proxies dict: Must have keys "http" and "https" and values of the proxies address:port to be used on the requests library format. You can also pass a list of proxies to rotate through. If an element of the list is None, it means not to use any proxy.
        :param debounce bool: Whether to try to debounce the requests or not. If set to False it will try to use proxies immediately upon failure.
        :param use_tor_fallback: If set to True, whenever a request fail will try to use a new tor session as a proxy. This will override the proxy options. Even when this is set to true it will still try to make a connection without using tor at first.
        """
        self.cookies = None
        self.context = context
        self.proxies = proxies
        self.debounce = debounce
        self.use_tor = use_tor_fallback

        if not isinstance(self.use_tor, bool):
            raise ValueError("use_tor_fallback must be a boolean")

        if proxies is None:
            self.proxies = [None]
        elif isinstance(proxies, dict):
            self.proxies = [proxies]
        elif isinstance(proxies, list):
            self.proxies = []
            for proxy in proxies:
                if isinstance(proxy, dict):
                    self.proxies.append(proxy)
                elif isinstance(proxy, str):
                    self.proxies.append({"http": proxy, "https": proxy})
                elif proxy is None:
                    self.proxies.append(None)
                else:
                    raise ValueError(
                        "Invalid proxy format. Must be a dict or a list of dicts."
                    )
        else:
            raise ValueError(
                "Invalid proxy format. Must be a dict, a list of dicts or a list of strings."
            )

        self._proxy_index = 0
        self._debounce_attempts = 0
        self._proxy = self.proxies[0]
        self._attempts = 0
        self.session = None
        self._tor_context = None
        self._start_session()

    def _start_session(self, tor=False):
        """Starts a new session refreshing the cookies"""
        logging.debug("Starting new session")
        if tor:
            try:
                from torpy.http.requests import tor_requests_session
            except ModuleNotFoundError:
                raise ModuleNotFoundError("You have to install torpy to use the tor fallback. Run: pip3 install torpy[requests]")
            logging.info("Starting tor session")
            if self._tor_context:
                self._tor_context.__exit__()
            self._tor_context = tor_requests_session()
            self.session = self._tor_context.__enter__()
            logging.info("tor session created!")
        else:
            self.session = requests.Session()
            self.session.proxies = self._proxy

        date = datetime.now().strftime("%Y%m%d")
        response = self.session.get(CLEVERBOT_COOKIE_URL + date)
        if "Set-cookie" not in response.headers:
            if not tor and self.use_tor:
                return self._start_session(self.use_tor)
            raise requests.exceptions.HTTPError(f"Api returned '{response.status_code}' for request. Try to use a proxy or use_tor_fallback")
        self.cookies = {
            "XVIS": re.search(r"\w+(?=;)", response.headers["Set-cookie"]).group()
        }

    def send(self, message: str) -> str:
        """
        Sends a message to Cleverbot and returns the response as a string.

        :param message str: The message to send to Cleverbot.
        :rtype str: Message received from Cleverbot.
        """
        payload = f"stimulus={requests.utils.requote_uri(message)}&"
        _context = self.context[:]
        reverse_context = list(reversed(_context))

        # Clear context list to keep things short and fast
        self._clear_context()
        for i in range(len(_context)):
            payload += f"vText{i + 2}={requests.utils.requote_uri(reverse_context[i])}&"

        # Append message to the context for future messages
        self.context.append(message)

        payload += "cb_settings_scripting=no&islearning=1&icognoid=wsf&icognocheck="

        # Checksum
        payload += hashlib.md5(payload[7:33].encode()).hexdigest()

        response = self.session.post(
            CLEVERBOT_API_URL,
            cookies=self.cookies,
            data=payload,
        )

        # If the request is not succesful, refresh cookies, debounce, try a new proxy
        if response.status_code != 200:
            if self._attempts > len(self.proxies) + MAX_ATTEPTS:
                logging.error(
                    f"Cleverbot failed to respond after {self._attempts} attempts. Giving up."
                )
                self._attempts = 0
                # TODO: Raise exception? Return None? idk
                return ""

            self._start_session(self.use_tor)
            if self._debounce_attempts >= MAX_DEBOUNCE_ATTEMPS or not self.debounce:
                self._debounce_attempts = 0
                self._proxy_index += 1
                if self._proxy_index >= len(self.proxies):
                    self._proxy_index = 0
                logging.info(
                    f"Cleverbot failed to respond. Trying proxy {self.proxies[self._proxy_index]}"
                )
                self._proxy = self.proxies[self._proxy_index]

            elif self.debounce:
                self._debounce_attempts += 1
                logging.info(
                    f"Cleverbot failed to respond. Trying again in {DEBOUNCE_TIME * self._debounce_attempts} seconds."
                )
                time.sleep(DEBOUNCE_TIME * self._debounce_attempts)

            self.context.pop()
            self._attempts += 1
            return self.send(message)

        response = re.split(r"\r", response.content.decode("utf-8"))[0]
        self._attempts = 0

        # Append bot's response to the context
        self.context.append(response)
        return codecs.escape_decode(bytes(response, "utf-8"))[0].decode("utf-8")

    def _clear_context(self):
        """_clear_context."""
        if len(self.context) >= MAX_CONTEXT_LENGTH:
            self.context.pop(0)


from unittest import TestCase

import pytest

from environments.identities.models import Identity
from environments.models import Environment
from organisations.models import Organisation
from projects.models import Project
from segments.models import PERCENTAGE_SPLIT, Condition, Segment, SegmentRule

@pytest.mark.django_db
class SegmentRuleTest(TestCase):
    def setUp(self) -> None:
        self.organisation = Organisation.objects.create(name="Test Org")
        self.project = Project.objects.create(
            name="Test Project", organisation=self.organisation
        )
        self.environment = Environment.objects.create(
            name="Test Environment", project=self.project
        )
        self.identity = Identity.objects.create(
            environment=self.environment, identifier="test_identity"
        )
        self.segment = Segment.objects.create(project=self.project, name="test_segment")

    def test_get_segment_returns_parent_segment_for_nested_rule(self):
        # Given
        parent_rule = SegmentRule.objects.create(
            segment=self.segment, type=SegmentRule.ALL_RULE
        )
        child_rule = SegmentRule.objects.create(
            rule=parent_rule, type=SegmentRule.ALL_RULE
        )
        grandchild_rule = SegmentRule.objects.create(
            rule=child_rule, type=SegmentRule.ALL_RULE
        )
        Condition.objects.create(
            operator=PERCENTAGE_SPLIT, value=0.1, rule=grandchild_rule
        )

        # When
        segment = grandchild_rule.get_segment()

        # Then
        assert segment == self.segment


import torch
from torch.nn import LSTM, Embedding, Linear, Module, Dropout, LeakyReLU, GRU, LogSoftmax
import json
from torch.optim import SGD, Adam
OPT = {"Adam": Adam, "SGD": SGD}

class MTEncoder(Module):
    """
    this class implements Sentence Encoder
    It's full operation consists of:
        1. get sentences (words are indices)
        2. obtain Embeddings
        3. pass through on directional LSTM
        4. return last layer
    """
    def __init__(self, params, vocab_dim, pad_index):
        super().__init__()
        self._params = params["encoder"] if type(params) is dict else json.load(open(params, "rt"))["encoder"]
        self._embeddings = Embedding(vocab_dim, self._params["embed_dim"], padding_idx=pad_index)   # word embed layer
        self._gru = GRU(self._params["embed_dim"], self._params["hidden_dim"],               # LSTM
                        self._params["num_layers"], dropout=self._params["dropout"], batch_first=True,
                        bidirectional=True)
        self._dropout = Dropout(p=self._params["dropout"])

    def forward(self, words_embed):
        """
        :param words_embed: word indices batch
        :return: sentence representation vector
        """
        x = self._dropout(self._embeddings(words_embed))
        output_seq, hidden_seq = self._gru(x)
        return output_seq

class MTDecoder(Module):
    """
    this class implements Sentence Decoder - it generates a new sentence given a context vector
    It's full operation consists of:
        1. get context vector
        2. init hidden vector
        3. loop
            3.1 predict next word using the hidden and context vectors
            3.2 update hidden vector
            3.3 save prediction and continue
    """
    def __init__(self, params, vocab_dim, start_index, context_dim):
        super().__init__()
        self._params = params["decoder"] if type(params) is dict else json.load(open(params, "rt"))["decoder"]

        self._embeddings = Embedding(vocab_dim, self._params["embed_dim"])                          # word embed layer
        self._lstm = LSTM(self._params["embed_dim"] + context_dim,                                  # LSTM layers
                          self._params["hidden_dim"], self._params["num_layers"],
                          dropout=self._params["dropout"], batch_first=True)
        self._out_dim = vocab_dim
        self._linear = Linear(self._params["hidden_dim"], vocab_dim)                           # Linear layer
        self._start = start_index
        self._dropout = Dropout(p=self._params["dropout"])
        self._softmax = LogSoftmax(dim=2)

    def forward(self, context, max_len, target=None, criterion=None):
        """
        :param context: context vector batch
        :param max_len: the function will create sequences of length max_len
        :return: generated sequences
        """
        context = self._dropout(context)
        x = torch.ones((context.shape[0], 1)).long() * self._start                     # init start vector
        hidden = None                                                                  # init hidden vector, and out
        out = torch.zeros((context.shape[0], max_len, self._out_dim))
        loss = 0
        for i in range(max_len):
            ex = self._embeddings(x)                                                   # obtain embeddings of last word
            # concat last word and context vectors and pass through LSTM
            lstm_input = torch.cat((ex, context.unsqueeze(dim=1)), dim=2)
            output_seq, hidden = self._lstm(lstm_input, hidden)
            word_distribution = self._softmax(self._linear(output_seq))         # predict next word distribution

            loss += criterion(word_distribution, target[:, i].unsqueeze(dim=0)) if criterion is not None else 0

            out[:, i, :] = word_distribution.squeeze(dim=1)
            x = word_distribution.argmax(dim=2).detach() if target is None else target[:, i].unsqueeze(dim=1)
        return out, loss

class MTEncoderDecoder(Module):
    def __init__(self, params, start_index, sourcee_pad, source_vocab_dim, target_vocab_dim):
        """
        this class implements a Full Machine Translation Model based on Encoder-Decoder architecture
        It's full operation consists of:
            1. Encode - obtain context vector for source sentence
            2. Decode = generate new sentence by the context vector
        """
        # super(MTEncoderDecoder, self).__init__()
        super().__init__()
        self._params = params if type(params) is dict else json.load(open(params, "rt"))
        self._encoder = MTEncoder(params, source_vocab_dim, sourcee_pad)
        self._decoder = MTDecoder(params, target_vocab_dim, start_index, 4 * self._params["encoder"]["hidden_dim"])
        # self.set_optimizer(self._params["optimizer"])

    def set_optimizer(self, opt_params):
        return OPT[opt_params["type"]](self.parameters(), **opt_params["kwargs"])

    def forward(self, word_embed, target=None, criterion=None):
        """
        :param word_embed: bach of source sentences
        :return: new generated sentence (word distribution)
        """

        context = self._encoder(word_embed)
        context = torch.cat([context[:, 0, :], context[:, -1, :]], dim=1)  # first and last hidden
        max_len = 2 * word_embed.shape[1] if target is None else word_embed.shape[1] - 1
        generated, loss = self._decoder(context, max_len=max_len, target=target, criterion=criterion)
        return generated, loss

def sanity_test():
    from dataset import MTDataset
    from torch.utils.data import DataLoader
    train = MTDataset("../data/train.src", "../data/train.trg")
    model = MTEncoderDecoder("sample_params.json", train.target_start_index, train.source_pad_index,
                             train.len_source_vocab, train.len_target_vocab)

    dl = DataLoader(train,
                    num_workers=4,
                    shuffle=True,
                    batch_size=64,
                    collate_fn=train.collate_fn)

    for src_, dst_ in dl:
        print(model(src_, target=dst_))

if __name__ == '__main__':
    sanity_test()


import sys
import random
import pygame
from pygame.locals import *
from Player import Player
from card_data import card_deck

DEFAULT_CARD_SIZE = (100,145)
PLAYER_IMAGE_POS = (100,180)
ENEMY_IMAGE_POS = (485,180)
PLAYER_NUMBER_POS = (50,330)
ENEMY_NUMBER_POS = (430,330)
RESULT_POS = (250,200)
WAR_POS = (300,200)
START_POSE = (0,420)
PLAYER_NAME = (100,110)
ENEMY_NAME = (485,110)
BATTLE_NUMBER_POS = (230,30)
IMAGE_SIZE = (700,500)
WAR_PHASE_POS = (150,430)

class WarCardGame:

    def __init__(self):
        self.screen = pygame.display.set_mode((700,500))
        self.screen.fill((0, 102, 0))
        pygame.display.set_caption("Card War Game")
        pygame_icon = pygame.image.load("cards/ace_of_spades.png")
        pygame.display.set_icon(pygame_icon)

    def run_game(self):
        pygame.init()
        global player, enemy, war_award, war_phase
        war_phase = False
        war_award = []
        player = Player()
        enemy = Player()
        self.start_screen()
        self.shuffle_cards()

        '''Game Loop'''
        is_result = False

        number_of_battles = 1
        clock = pygame.time.Clock()

        while True:
            for event in pygame.event.get():

                if event.type ==  pygame.QUIT:
                    sys.exit()

                if event.type == pygame.KEYDOWN:

                    if event.key == K_ESCAPE:
                        sys.exit()

                    elif event.key != K_SPACE :
                        continue

                    else:
                        if event.key == K_SPACE:
                            self.screen.fill((0, 102, 0))
                            self.draw_info(number_of_battles)
                            if not is_result:
                                number_of_battles+=1
                                result = self.draw_cards()
                                self.draw_result(result)
                                is_result = True
                            else:
                                self.draw_cards_backwards()
                                self.draw_number_of_cards()
                                is_result = False

                        if self.end_of_game(result):
                            break

            pygame.display.flip()
            clock.tick(30)

    def shuffle_cards(self):
        '''Shuffle cards at the beginning of the game'''
        full_deck = random.sample(card_deck.keys(),52)
        player.deck = full_deck[:26]
        enemy.deck = full_deck[26:]

    def change_decks(self,player_card,enemy_card,result):
        '''Adding cards to winner's deck'''
        global war_phase,war_award

        if war_phase:
            self.draw_cards_backwards()
            font = pygame.font.SysFont("comicsansms", 42)
            player_number_text = font.render(f"BLIND WAR PHASE",False,(255, 204, 102))
            self.screen.blit(player_number_text,WAR_PHASE_POS)

        if result == "war":
            if not war_award:
                self.war(player_card,enemy_card)
                war_phase = True

        elif result == "win" and war_phase == False:
            card = enemy.deck.pop(enemy_card)
            player.deck.append(card)
            if war_award:
                enemy.deck.extend(war_award)
                war_award.clear()

        elif result == "lose" and war_phase == False:
            card = player.deck.pop(player_card)
            enemy.deck.append(card)
            if war_award:
                enemy.deck.extend(war_award)
                war_award.clear()

        elif war_phase == True:
            self.war(player_card,enemy_card)
            war_phase = False

        player.number_of_cards = len(player.deck)
        enemy.number_of_cards = len(enemy.deck)

    def war(self,player_card,enemy_card):
            '''Adding cards to the war reward array'''
            card = enemy.deck.pop(enemy_card)
            war_award.append(card)
            card = player.deck.pop(player_card)
            war_award.append(card)

    def compare_cards(self,player_choice,enemy_choice):
        '''Compare value of each chosen card'''
        player_value = card_deck[player.deck[player_choice]]["value"]
        enemy_value = card_deck[enemy.deck[enemy_choice]]["value"]
        if player_value > enemy_value:
            return "win"
        elif player_value == enemy_value:
            return "war"
        else:
            return "lose"

    def end_of_game(self,result):
        '''Drawing the result of the game'''
        if player.number_of_cards == 0 or (result == "war" and player.number_of_cards < 2):
            self.screen.fill((0, 102, 0))
            game_result = pygame.image.load("cards\you_lost.png")
            game_result = pygame.transform.scale(game_result,IMAGE_SIZE)
            self.screen.blit(game_result,(0,0))
            return True
        elif enemy.number_of_cards == 0 or (result == "war" and enemy.number_of_cards < 2):
            self.screen.fill((0, 102, 0))
            game_result = pygame.image.load("cards\game_won.png")
            game_result = pygame.transform.scale(game_result,IMAGE_SIZE)
            self.screen.blit(game_result,(0,0))
        else:
            return False

    def start_screen(self):
        '''Drawing screen at the beginning'''
        font = pygame.font.SysFont("comicsansms", 42)
        player_number_text = font.render(f"To start press SPACEBAR",False,(0,0,0))
        self.screen.blit(player_number_text,START_POSE)

    def draw_cards(self):
        '''Get number of cards in each deck, pick a random card from the deck and draw the card on the surface. Getting result of the battle and modify decks'''

        player.number_of_cards = len(player.deck)
        player_choice = random.randint(0,player.number_of_cards-1)
        image_path = card_deck[player.deck[player_choice]]["image"]
        card = pygame.image.load(image_path)
        card = pygame.transform.scale(card,DEFAULT_CARD_SIZE)
        self.screen.blit(card,PLAYER_IMAGE_POS)

        enemy.number_of_cards = len(enemy.deck)
        enemy_choice = random.randint(0,enemy.number_of_cards-1)

        image_path = card_deck[enemy.deck[enemy_choice]]["image"]
        card = pygame.image.load(image_path)
        card = pygame.transform.scale(card,DEFAULT_CARD_SIZE)
        self.screen.blit(card,ENEMY_IMAGE_POS)

        result = self.compare_cards(player_choice,enemy_choice)
        self.change_decks(player_choice,enemy_choice,result)

        return result

    def draw_number_of_cards(self):
        '''Drawing number of cards which remain in each deck'''
        font = pygame.font.SysFont("comicsansms", 32)
        player_number_text = font.render(f"Cards left: {player.number_of_cards}",False,(0,0,0))
        self.screen.blit(player_number_text,PLAYER_NUMBER_POS)
        enemy_number_text = font.render(f"Cards left: {enemy.number_of_cards}",False,(0,0,0))
        self.screen.blit(enemy_number_text,ENEMY_NUMBER_POS)

    def draw_result(self,result):
        '''Drawing result of the battle'''
        font = pygame.font.SysFont("comicsansms", 32)
        if result == "win":
            result_text = font.render(f"YOU WON",False,(153, 255, 153))
            self.screen.blit(result_text,RESULT_POS)
        elif result == "war":
            result_text = font.render(f"WAR",False,(255, 204, 102))
            self.screen.blit(result_text,WAR_POS)
        else:
            result_text = font.render(f"YOU LOST",False,(255, 51, 0))
            self.screen.blit(result_text,RESULT_POS)

    def draw_info(self,number_of_battles):
        '''Drawing names and number of battles which were made'''

        font = pygame.font.SysFont("comicsansms", 32)
        text = font.render(f"Player",False,(0,0,0))
        self.screen.blit(text,PLAYER_NAME)
        text = font.render(f"Enemy",False,(0,0,0))
        self.screen.blit(text,ENEMY_NAME)
        text = font.render(f"Battle nr. {number_of_battles}",False,(0,0,0))
        self.screen.blit(text,BATTLE_NUMBER_POS)

    def draw_cards_backwards(self):
        '''Drawing back of the card before battle'''
        card = pygame.image.load("cards/card_back.png")
        card = pygame.transform.scale(card,DEFAULT_CARD_SIZE)
        self.screen.blit(card,ENEMY_IMAGE_POS)
        self.screen.blit(card,PLAYER_IMAGE_POS)

if __name__ == "__main__":
    game = WarCardGame()
    game.run_game()