import logging
import os.path
import re
from argparse import ArgumentParser
from collections import defaultdict
from dataclasses import dataclass
from typing import Any, Dict, Iterator, List, NamedTuple, Optional, Union, cast

from common import directories, file_utils
from common.commands.base import ArxivBatchCommand
from common.parse_tex import PhraseExtractor, overlaps
from common.types import ArxivId, CharacterRange, FileContents
from tqdm import tqdm
from typing_extensions import Literal

from ..nlp import DefinitionDetectionModel
from ..types import Definiendum, Definition, EmbellishedSentence, TermReference

"""
Deployment TODOs:
* TODO(dykang): document model download instructions in README
* TODO(dykang): make model location configurable in config.ini
"""

TermName = str
DefinitionId = str

@dataclass(frozen=True)
class DetectDefinitionsTask:
    arxiv_id: ArxivId
    sentences: List[EmbellishedSentence]
    tex_by_file: Dict[str, FileContents]

@dataclass(frozen=True)
class TermDefinitionPair:
    term_start: int
    term_end: int
    term_text: str
    term_type: str
    term_confidence: Optional[float]
    definition_start: int
    definition_end: int
    definition_text: str
    definition_type: str
    definition_confidence: Optional[float]

def get_token_character_ranges(text: str, tokens: List[str]) -> List[CharacterRange]:
    """
    Extract start and end charcter positions for each token in featurized tokens
    """
    ranges = []
    current_position = 0
    for token in tokens:
        start_index = text[current_position:].index(token)
        ranges.append(
            CharacterRange(
                current_position + start_index,
                current_position + start_index + len(token) - 1,
            )
        )
        current_position += len(token) + start_index
    return ranges

StringOffset = int

def get_symbol_texs(
    sentence_with_symbol_tags: str, sentence_with_formula_contents: str
) -> Optional[Dict[StringOffset, str]]:
    """
    Get a map from the character offsets of the word 'SYMBOL' in the first input string
    to the TeX for that symbol, found from the second input string.  Example:

    Input: get_symbol_texs("Add SYMBOL to SYMBOL.", "Add (((FORMULA:x))) to (((FORMULA:y)))")
    Output: { 4: "x", 14: "y" }

    Returns None if there was an error establishing a mapping, for instance if there are
    a different number of 'SYMBOL' tags between the two versions of the sentence.
    """
    symbol_starts = [
        match.start() for match in re.finditer(r"SYMBOL", sentence_with_symbol_tags)
    ]
    symbol_texs = [
        match.group(1)
        for match in re.finditer(
            # The extra '[)]*]' before the ending triple-parentheses ensures that
            # closing parentheses that were part of the original equation (e.g., for
            # a symbol like 'f(x)') are included in the symbol, rather than being read as part
            # of the triple-parens formula delimiter.
            r"\(\(\(FORMULA:(.*?[)]*)\)\)\)",
            sentence_with_formula_contents,
        )
    ]

    if len(symbol_starts) != len(symbol_texs):
        logging.warning(  # pylint: disable=logging-not-lazy
            "The two representations of a sentence %s and %s were detected as having differing "
            + "numbers of symbols. A lookup table for symbol TeX cannot be built. The TeX for "
            + "symbols in these equations may not be correct in the output.",
            sentence_with_symbol_tags,
            sentence_with_formula_contents,
        )
        return None

    return dict(zip(symbol_starts, symbol_texs))

def consolidate_keyword_definitions(
    text: str,
    tokens: List[str],
    slot_predictions: List[str],
    slot_prediction_confidences: List[float],
    prediction_type: str,
) -> List[TermDefinitionPair]:

    # Make index from tokens to their character positions.
    ranges = get_token_character_ranges(text, tokens)

    # Extract ranges for all terms and definitions.
    terms: List[List[CharacterRange]] = []
    definitions: List[List[CharacterRange]] = []
    term_confidence_list: List[float] = []
    definition_confidence_list: List[float] = []
    term_ranges, definition_ranges = [], []
    term_confidences, definition_confidences = [], []

    for (slot_label, slot_confidence, range_) in zip(
        slot_predictions, slot_prediction_confidences, ranges
    ):
        if slot_label == "TERM":
            term_ranges.append(range_)
            term_confidences.append(slot_confidence)
        if slot_label == "DEF":
            definition_ranges.append(range_)
            definition_confidences.append(slot_confidence)
        if slot_label == "O":
            if len(term_ranges) > 0:
                terms.append(term_ranges)
                term_confidence_list.append(
                    sum(term_confidences) / len(term_confidences)
                )
            if len(definition_ranges) > 0:
                definitions.append(definition_ranges)
                definition_confidence_list.append(
                    sum(definition_confidences) / len(definition_confidences)
                )
            term_ranges, definition_ranges = [], []

    if len(term_ranges) > 0:
        terms.append(term_ranges)
        term_confidence_list.append(sum(term_confidences) / len(term_confidences))
    if len(definition_ranges) > 0:
        definitions.append(definition_ranges)
        definition_confidence_list.append(
            sum(definition_confidences) / len(definition_confidences)
        )
    term_ranges, definition_ranges = [], []

    # Match pairs of term and definitions sequentially ( O T O O D then (T, D)). This
    # does not handle multi-term / definitions pairs.
    num_term_definition_pairs = min(len(terms), len(definitions))
    pairs: List[TermDefinitionPair] = []
    for pair_index in range(num_term_definition_pairs):
        logging.debug(
            "Found slot predictions for tokens: (tokens: %s), (slots: %s)",
            tokens,
            slot_predictions,
        )
        logging.debug(
            "Processing term-definition pair %d of %d",
            pair_index,
            num_term_definition_pairs,
        )

        term_rangelist = terms[pair_index]
        definition_rangelist = definitions[pair_index]
        term_confidence = term_confidence_list[pair_index]
        definition_confidence = definition_confidence_list[pair_index]

        term_start = min([term_range.start for term_range in term_rangelist])
        term_end = max([term_range.end for term_range in term_rangelist]) + 1
        definition_start = min(
            [definition_range.start for definition_range in definition_rangelist]
        )
        definition_end = (
            max([definition_range.end for definition_range in definition_rangelist]) + 1
        )

        all_definition_types = {
            "W00": ["term", "definition"],
            "AI2020": ["abbreviation", "expansion"],
            "DocDef2": ["symbol", "nickname"],
        }

        term_type = all_definition_types[prediction_type][0]
        definition_type = all_definition_types[prediction_type][1]

        pair = TermDefinitionPair(
            term_start=term_start,
            term_end=term_end,
            term_text=text[term_start:term_end],
            term_type=term_type,
            term_confidence=term_confidence,
            definition_start=definition_start,
            definition_end=definition_end,
            definition_text=text[definition_start:definition_end],
            definition_type=definition_type,
            definition_confidence=definition_confidence,
        )
        logging.debug("Found definition-term pair %s", pair)
        pairs.append(pair)

    return pairs

def check_text_contains_abbreviation_for_sanity(
    abbreviation_text: str,
    expansion_text: str,
    abbreviation_start: int,
    abbreviation_end: int,
    expansion_start: int,
    expansion_end: int,
) -> bool:
    """
    Very smooth filter:
    Outputs from our current abbreviation detector include lots of noise
        e.g., "alpha x (Ref )" or any phrases with parenthesis
    This function filters out the noise using simple heuristics.
    """
    # If abbreviation and expansion overlaps in positions, ignore them.
    if (
        len(
            set(range(abbreviation_start, abbreviation_end - 1)).intersection(
                range(expansion_start, expansion_end - 1)
            )
        )
        > 0
    ):
        return False

    # If citation patterns are detected (e.g., Citation (CITATION)), ignore them.
    if "citation" in abbreviation_text.lower():
        return False

    # Check each letter in abbreviation appears in expansion text.
    is_abbreviation = True
    for letter_in_abbreviation in abbreviation_text.lower():
        if letter_in_abbreviation not in expansion_text.lower():
            is_abbreviation = False

    return is_abbreviation

def get_abbreviations(
    text: str, tokens: List[str], nlp_model: Any
) -> List[TermDefinitionPair]:
    """
    Legacy function for detecting abbreviations in a text. No longer in use.
    """

    pairs: List[TermDefinitionPair] = []

    doc = nlp_model(text)

    # Make index from tokens to their character positions.
    token_ranges = get_token_character_ranges(text, tokens)

    for abbreviation in doc._.abbreviations:
        # Abbreviation (term).
        abbreviation_ranges = token_ranges[abbreviation.start : abbreviation.end]
        abbreviation_start = min([r.start for r in abbreviation_ranges])
        abbreviation_end = max([r.end for r in abbreviation_ranges]) + 1

        # Expansion (definition).
        expansion_ranges = token_ranges[
            abbreviation._.long_form.start : abbreviation._.long_form.end
        ]
        expansion_start = min(
            [expansion_range.start for expansion_range in expansion_ranges]
        )
        expansion_end = (
            max([expansion_range.end for expansion_range in expansion_ranges]) + 1
        )

        if not check_text_contains_abbreviation_for_sanity(
            str(abbreviation),
            str(abbreviation._.long_form),
            abbreviation_start,
            abbreviation_end,
            expansion_start,
            expansion_end,
        ):
            continue

        pair = TermDefinitionPair(
            term_start=abbreviation_start,
            term_end=abbreviation_end,
            term_text=text[abbreviation_start:abbreviation_end],
            term_type="abbreviation",
            term_confidence=None,
            definition_start=expansion_start,
            definition_end=expansion_end,
            definition_text=text[expansion_start:expansion_end],
            definition_type="expansion",
            definition_confidence=None,
        )
        logging.debug("Found abbreviation-expansion pair %s", pair)
        pairs.append(pair)

    return pairs

class Nickname(NamedTuple):
    token_indexes: List[int]
    token_ranges: List[CharacterRange]

Direction = Literal["LEFT", "RIGHT"]

def search_symbol_nickname(
    token_index: int,
    pos: List[str],
    ranges: List[CharacterRange],
    direction: Direction,
) -> Optional[Nickname]:
    """
    This function searches for a nickname pattern for a token, in the direction specified.
    Get all tokens to the left or right of the token that match the allowed POS tags.
    Return symbol and nickname indices and ranges if there are valid POS spans. Else, return None for indices and spans.
    """

    # Union set of POS tags in nicknames. Feel free to add more if you have new patterns.
    # POS Tag Expansions : https://www.ling.upenn.edu/courses/Fall_2003/ling001/penn_treebank_pos.html.
    UNION_POS_LIST = ["DT", "JJ", "NN", "NNS", "NNP", "NNPS", ":", "-LRB-"]
    pos_indexes = []
    pos_tags = []
    pos_ranges = []
    if direction == "RIGHT":
        current_index = token_index + 1
    else:
        current_index = token_index - 1

    # Exit if the current index is greater than the length of the sentence or below zero.
    if current_index >= len(pos) or current_index < 0:
        return None

    # Search ahead until the POS tag is not in the union set.
    while pos[current_index] in UNION_POS_LIST:
        pos_indexes.append(current_index)
        pos_ranges.append(ranges[current_index])
        pos_tags.append(pos[current_index])
        if direction == "RIGHT":
            current_index += 1
            if current_index >= len(pos):
                break
        else:
            current_index -= 1
            if current_index < 0:
                break

    # Ignore when the POS tag is among ["DT", ":", "-LRB-"].
    if len(pos_tags) == 1 and (pos_tags[0] in ["DT", ":", "-LRB-"]):
        return None

    if len(pos_indexes) > 0:
        if direction == "LEFT":
            nickname_indexes = list(reversed(pos_indexes))
            nickname_ranges = list(reversed(pos_ranges))
            nickname_tags = list(reversed(pos_tags))
        else:
            nickname_indexes = pos_indexes
            nickname_ranges = list(pos_ranges)
            nickname_tags = list(pos_tags)

        # Skip ["DT", ":", "-LRB-"] as first or last nickname tokens.
        if nickname_tags[0] in ["DT", ":", "-LRB-"]:
            nickname_ranges = nickname_ranges[1:]
            nickname_indexes = nickname_indexes[1:]
        elif nickname_tags[-1] in ["DT", ":", "-LRB-"]:
            nickname_ranges = nickname_ranges[:-1]
            nickname_indexes = nickname_indexes[:-1]

        return Nickname(nickname_indexes, nickname_ranges)
    else:
        return None

class SymbolNickname(NamedTuple):
    symbol: str
    symbol_index: int
    symbol_range: CharacterRange
    nickname_indexes: List[int]
    nickname_ranges: List[CharacterRange]

def get_symbol_nickname_pairs(
    text: str, tokens: List[str], pos: List[str], symbol_texs: Dict[StringOffset, str]
) -> List[TermDefinitionPair]:
    """
    Legacy function for detecting nicknames of a symbol in a text. No longer in use.
    """

    # Check whether a symbol's definition is a nickname of the symbol or not
    # using heuristic rules below, although they are not perfect for some cases.
    # a/DT particular/JJ transcript/NN SYMBOL/NN.
    # the/DT self-interaction/JJ term/JJ SYMBOL/NN.
    # the/DT main/JJ effect/NN of/NN feature/JJ SYMBOL/JJ.

    # If a new algorithm needs to be added, a separate function can be called in the iteration below. Finally, append a pair object to the 'pairs' list.
    ranges = get_token_character_ranges(text, tokens)
    pairs = []
    if "SYMBOL" in text:

        symbol_nickname_pairs: List[SymbolNickname] = []
        symbol_index = 0

        for token_index, (token, range_) in enumerate(zip(tokens, ranges)):

            # 1. If of the form '*th', check RIGHT of symbol.
            # 2. If token is a symbol, then:
            #   a. If the symbol tex is present:
            #       i. If single length symbol, first check LEFT then RIGHT.
            #       ii. If multi length symbol, check LEFT.
            #   b. If symbol tex is not present, just check LEFT.
            if token == "SYMBOLth":
                nickname = search_symbol_nickname(token_index, pos, ranges, "RIGHT")
                if nickname is not None and range_.start in symbol_texs:
                    symbol_nickname_pairs.append(
                        SymbolNickname(
                            symbol_texs[range_.start],
                            token_index,
                            range_,
                            nickname.token_indexes,
                            nickname.token_ranges,
                        )
                    )

            elif token == "SYMBOL":
                # Decide the order of LEFT or RIGHT.
                directions: List[Direction] = []
                if token_index > 0:
                    if token_index < len(pos) - 1:
                        # Search RIGHT first if the immediate next token is among ["NN", ":"].
                        if pos[token_index + 1] in ["NN", ":"]:
                            directions = ["RIGHT", "LEFT"]
                        else:
                            directions = ["LEFT", "RIGHT"]
                    else:
                        directions = ["LEFT"]
                elif token_index < len(pos) - 1:
                    directions = ["RIGHT"]

                for direction in directions:
                    nickname = search_symbol_nickname(
                        token_index, pos, ranges, direction
                    )
                    if nickname is not None and range_.start in symbol_texs:
                        symbol_nickname_pairs.append(
                            SymbolNickname(
                                symbol_texs[range_.start],
                                token_index,
                                range_,
                                nickname.token_indexes,
                                nickname.token_ranges,
                            )
                        )
                        break

            if "SYMBOL" in token:
                symbol_index += 1

        for sn in symbol_nickname_pairs:
            symbol_start = sn.symbol_range.start
            symbol_end = sn.symbol_range.end + 1
            nickname_start = min(
                [nickname_range.start for nickname_range in sn.nickname_ranges]
            )
            nickname_end = (
                max([nickname_range.end for nickname_range in sn.nickname_ranges]) + 1
            )
            nickname_text = text[nickname_start:nickname_end]
            # Reject a nickname if the nickname text is "SYMBOL".
            if nickname_text == "SYMBOL":
                continue

            pair = TermDefinitionPair(
                term_start=symbol_start,
                term_end=symbol_end,
                term_text=sn.symbol,
                term_type="symbol",
                term_confidence=None,
                definition_start=nickname_start,
                definition_end=nickname_end,
                definition_text=nickname_text,
                definition_type="nickname",
                definition_confidence=None,
            )
            logging.debug("Found symbol-nickname pair %s", pair)
            pairs.append(pair)

    return pairs

class DetectDefinitions(
    ArxivBatchCommand[
        DetectDefinitionsTask, Union[Definiendum, Definition, TermReference]
    ]
):
    """
    Extract terms and definitions of those terms from sentences using a pre-trained definition
    extraction model. Performs the following steps:
    1. Load cleaned sentences from the 'detected-sentences' data directory
    2. Extract features from each sentence (aka featurization)
    3. Load the pre-trained NLP model, load the features, and predict intent and term:definition
       slots. The model predicts intents and slots:
        * Intent: whether a sentence is detected as including a definition or not
        * Slot: a tag for each token indicating it is one of (TERM, DEFINITION, neither)
    4. Find all references to terms in the text
    5. Save detected terms to 'entities-terms.csv' and detected definitions to
       'entities-definitions.csv'
    """

    @staticmethod
    def get_name() -> str:
        return "detect-definitions"

    @staticmethod
    def get_description() -> str:
        return "Extract term and definition pairs from LaTeX."

    @staticmethod
    def init_parser(parser: ArgumentParser) -> None:
        super(DetectDefinitions, DetectDefinitions).init_parser(parser)
        parser.add_argument(
            "--show-progress",
            action="store_true",
            help=(
                "Whether to show progress bar during definition extraction. "
                + "Interferes with log messages shown on the command line."
            ),
        )
        parser.add_argument(
            "--batch-size",
            type=int,
            default=8,
            help=("Number of sentences to process at a time to detect definitions."),
        )

    def get_arxiv_ids_dirkey(self) -> str:
        return "detected-sentences"

    def load(self) -> Iterator[DetectDefinitionsTask]:
        for arxiv_id in self.arxiv_ids:

            output_dir = directories.arxiv_subdir("detected-definitions", arxiv_id)
            file_utils.clean_directory(output_dir)

            # Load cleaned sentences for definition detection.
            detected_sentences_path = os.path.join(
                directories.arxiv_subdir("sentence-tokens", arxiv_id), "sentences.csv",
            )
            try:
                sentences = list(
                    file_utils.load_from_csv(
                        detected_sentences_path, EmbellishedSentence
                    )
                )
            except FileNotFoundError:
                logging.warning(  # pylint: disable=logging-not-lazy
                    "No sentences data found for arXiv paper %s. Try re-running the pipeline, "
                    + "this time enabling the processing of sentences. If that doesn't work, "
                    + "there is likely an error in detected sentences for this paper.",
                    arxiv_id,
                )
                continue

            # Read in all TeX. Once definition detection is finished, all the TeX will be searched
            # for references to the defined terms.
            tex_by_file = file_utils.read_tex(arxiv_id)

            yield DetectDefinitionsTask(arxiv_id, sentences, tex_by_file)

    def process(
        self, item: DetectDefinitionsTask
    ) -> Iterator[Union[Definiendum, Definition, TermReference]]:
        sentences_ordered = sorted(item.sentences, key=lambda s: s.start)
        num_sentences = len(sentences_ordered)
        end_position_of_last_sentence = sentences_ordered[-1].end

        if len(item.sentences) == 0:
            logging.warning(  # pylint: disable=logging-not-lazy
                "No sentences found for arXiv ID %s. Skipping detection of sentences "
                + "that contain entities.",
                item.arxiv_id,
            )
            return

        # Load the pre-trained definition detection model.
        prediction_type = "DocDef2+AI2020+W00"
        model = DefinitionDetectionModel(prediction_type)

        definition_index = 0
        features = []
        sentences: List[EmbellishedSentence] = []

        definiendums: Dict[TermName, List[Definiendum]] = defaultdict(list)
        term_phrases: List[str] = []
        abbreviations: List[str] = []
        symbol_nicks: List[str] = []
        definitions: Dict[DefinitionId, Definition] = {}

        with tqdm(
            total=num_sentences, disable=(not self.args.show_progress)
        ) as progress:

            for sentence_index, sentence in enumerate(sentences_ordered):
                progress.update(1)

                # Only attempt to process sentences that have been marked as likely to be proper
                # plaintext. Note that this means some sentences may be skipped that didn't pass
                # heuristics in the sentence extractor.
                if not sentence.validity_guess:
                    continue

                # Extract features from raw text.
                featurized_text = model.featurize(sentence.legacy_definition_input)
                features.append(featurized_text)
                sentences.append(sentence)

                # Process sentences in batches.
                if (
                    len(features) >= self.args.batch_size
                    or sentence_index == num_sentences - 1
                ):

                    # Detect terms and definitions in each sentence with a pre-trained definition
                    # extraction model, from the featurized text.

                    (_, slots, slots_confidence) = model.predict_batch(
                        cast(List[Dict[Any, Any]], features)
                    )

                    # Package extracted terms and definitions into a representation that's
                    # easier to process.
                    for (
                        s,
                        sentence_features,
                        termdef_sentence_slots,
                        termdef_sentence_slots_confidence,
                        abbrexp_sentence_slots,
                        abbrexp_sentence_slots_confidence,
                        symnick_sentence_slots,
                        symnick_sentence_slots_confidence,
                    ) in zip(
                        sentences,
                        features,
                        slots["W00"],
                        slots_confidence["W00"],
                        slots["AI2020"],
                        slots_confidence["AI2020"],
                        slots["DocDef2"],
                        slots_confidence["DocDef2"],
                    ):
                        # Extract TeX for each symbol from a parallel representation of the
                        # sentence, so that the TeX for symbols can be saved.
                        # Types of [term and definition] pairs.
                        #   [nickname and definition] for symbols.
                        #   [abbreviation and expansion] for abbreviations.
                        #   [term and definition] for other types.
                        symbol_texs = get_symbol_texs(
                            s.legacy_definition_input, s.with_formulas_marked
                        )

                        # Only process slots when they include both 'TERM' and 'DEFINITION'.
                        if (
                            "TERM" not in termdef_sentence_slots
                            or "DEF" not in termdef_sentence_slots
                        ):
                            term_definition_pairs = []
                        else:
                            term_definition_pairs = consolidate_keyword_definitions(
                                s.legacy_definition_input,
                                sentence_features["tokens"],
                                termdef_sentence_slots,
                                termdef_sentence_slots_confidence,
                                "W00",
                            )

                        if (
                            "TERM" not in abbrexp_sentence_slots
                            or "DEF" not in abbrexp_sentence_slots
                        ):
                            abbreviation_expansion_pairs = []
                        else:
                            abbreviation_expansion_pairs = consolidate_keyword_definitions(
                                s.legacy_definition_input,
                                sentence_features["tokens"],
                                abbrexp_sentence_slots,
                                abbrexp_sentence_slots_confidence,
                                "AI2020",
                            )

                        if (
                            "TERM" not in symnick_sentence_slots
                            or "DEF" not in symnick_sentence_slots
                        ):
                            symbol_nickname_pairs = []
                        else:
                            symbol_nickname_pairs = consolidate_keyword_definitions(
                                s.legacy_definition_input,
                                sentence_features["tokens"],
                                symnick_sentence_slots,
                                symnick_sentence_slots_confidence,
                                "DocDef2",
                            )

                        pairs = (
                            term_definition_pairs
                            + symbol_nickname_pairs
                            + abbreviation_expansion_pairs
                        )
                        for pair in pairs:
                            tex_path = s.tex_path
                            definiendum_id = (
                                f"definiendum-{tex_path}-{definition_index}"
                            )
                            definition_id = f"definition-{tex_path}-{definition_index}"
                            definiendum_text = pair.term_text
                            definiendum_type = pair.term_type
                            definition_type = pair.definition_type

                            definiendum_confidence = pair.term_confidence
                            definition_confidence = pair.definition_confidence

                            # Map definiendum and definition start and end positions back to
                            # their original positions in the TeX.
                            offsets = s.legacy_definition_input_journal.initial_offsets(
                                pair.term_start, pair.term_end
                            )
                            if offsets[0] is None or offsets[1] is None:
                                logging.warning(  # pylint: disable=logging-not-lazy
                                    "Could not find offsets of definiendum %s in original TeX "
                                    + "(from sentence %s, file %s, arXiv ID %s). Definiendum will not be saved.",
                                    pair.term_text,
                                    s.id_,
                                    s.tex_path,
                                    item.arxiv_id,
                                )
                                continue
                            definiendum_start = s.start + offsets[0]
                            definiendum_end = s.start + offsets[1]

                            offsets = s.legacy_definition_input_journal.initial_offsets(
                                pair.definition_start, pair.definition_end
                            )
                            if offsets[0] is None or offsets[1] is None:
                                logging.warning(  # pylint: disable=logging-not-lazy
                                    "Could not find offsets of definition %s in original TeX "
                                    + "(from sentence %s, file %s, arXiv ID %s). Definiendum will not be saved.",
                                    pair.definition_text,
                                    s.id_,
                                    s.tex_path,
                                    item.arxiv_id,
                                )
                                continue
                            definition_start = s.start + offsets[0]
                            definition_end = s.start + offsets[1]

                            # Extract document-level features from sentence.
                            position_ratio = (
                                definiendum_start / end_position_of_last_sentence
                            )
                            section_name = s.section_name

                            try:
                                tex = item.tex_by_file[tex_path]
                            except KeyError:
                                logging.warning(  # pylint: disable=logging-not-lazy
                                    "Could not find TeX for %s. TeX will not be included in "
                                    + "the output data for definition '%s' for term '%s'",
                                    tex_path,
                                    pair.definition_text,
                                    definiendum_text,
                                )
                                definiendum_tex = "NOT AVAILABLE"
                                definition_tex = "NOT AVAILABLE"
                            else:
                                if (
                                    definiendum_type == "symbol"
                                    and symbol_texs is not None
                                    and pair.term_start in symbol_texs
                                ):
                                    definiendum_tex = symbol_texs[pair.term_start]
                                    definiendum_text = definiendum_tex
                                else:
                                    definiendum_tex = tex.contents[
                                        definiendum_start:definiendum_end
                                    ]
                                definition_tex = tex.contents[
                                    definition_start:definition_end
                                ]

                            # Save the definition to file.
                            definition = Definition(
                                id_=definition_id,
                                start=definition_start,
                                end=definition_end,
                                definiendum=definiendum_text,
                                type_=definition_type,
                                tex_path=tex_path,
                                tex=definition_tex,
                                text=pair.definition_text,
                                context_tex=sentence.context_tex,
                                sentence_id=sentence.id_,
                                intent=True,
                                confidence=definition_confidence,
                            )
                            definitions[definition_id] = definition
                            yield definition

                            # Don't save the definiendum to file yet. Save it in memory first, and then
                            # save it to file once it's done being processed. It will need
                            # to be associated with other definitions. Also, other references
                            # to the term will be detected before this method is over.
                            definiendum = Definiendum(
                                id_=definiendum_id,
                                text=definiendum_text,
                                type_=definiendum_type,
                                confidence=definiendum_confidence,
                                # Link the definiendum to the text that defined it.
                                definition_id=definition_id,
                                # Because a term can be defined multiple places in the paper, these
                                # three lists of definition data will be filled out once all of the
                                # definitions have been found.
                                definition_ids=[],
                                definitions=[],
                                definition_texs=[],
                                sources=[],
                                start=definiendum_start,
                                end=definiendum_end,
                                tex_path=tex_path,
                                tex=definiendum_tex,
                                context_tex=sentence.context_tex,
                                sentence_id=sentence.id_,
                                # Document-level features below.
                                position_ratio=position_ratio,
                                position_ratios=[],
                                section_name=section_name,
                                section_names=[],
                            )
                            definiendums[definiendum_text].append(definiendum)
                            if definiendum.type_ == "term":
                                term_phrases.append(definiendum.text)
                            if definiendum.type_ == "abbreviation":
                                abbreviations.append(definiendum.text)
                            if definiendum.type_ == "symbol":
                                symbol_nicks.append(definiendum.text)

                            definition_index += 1

                    features = []
                    sentences = []

        logging.debug(  # pylint: disable=logging-not-lazy
            "Finished detecting definitions for paper %s. Now finding references to defined terms.",
            item.arxiv_id,
        )

        all_definiendums: List[Definiendum] = []
        for _, definiendum_list in definiendums.items():
            all_definiendums.extend(definiendum_list)

        definition_ids: Dict[TermName, List[DefinitionId]] = {}
        definition_texs: Dict[TermName, List[str]] = {}
        definition_texts: Dict[TermName, List[str]] = {}
        sources: Dict[TermName, List[str]] = {}
        position_ratios: Dict[TermName, List[float]] = {}
        section_names: Dict[TermName, List[str]] = {}

        # Associate terms with all definitions that apply to them.
        for term, definiendum_list in definiendums.items():
            definition_ids[term] = [
                definiendum.definition_id for definiendum in definiendum_list
            ]
            definition_texs[term] = [
                definitions[definiendum.definition_id].tex
                for definiendum in definiendum_list
            ]
            definition_texts[term] = [
                definitions[definiendum.definition_id].text
                for definiendum in definiendum_list
            ]
            sources[term] = ["model"] * len(definition_ids[term])
            position_ratios[term] = [
                definiendum.position_ratio for definiendum in definiendum_list
            ]
            section_names[term] = [
                definiendum.section_name
                for definiendum in definiendum_list
                if definiendum.section_name is not None
            ]

        # Associate each definiendum with all applicable definitions, and save them to file.
        for _, definiendum_list in definiendums.items():
            for definiendum in definiendum_list:
                definiendum.definition_ids.extend(definition_ids[definiendum.text])
                definiendum.definition_texs.extend(definition_texs[definiendum.text])
                definiendum.definitions.extend(definition_texts[definiendum.text])
                definiendum.sources.extend(sources[definiendum.text])
                definiendum.position_ratios.extend(position_ratios[definiendum.text])
                definiendum.section_names.extend(section_names[definiendum.text])
                yield definiendum

        # Detect all other references to the defined terms. Detect references to textual
        # terms and abbreviations. References to symbols need not be found here; they
        # will be detected automatically in the symbol extraction code.
        term_index = 0

        for tex_path, file_contents in item.tex_by_file.items():
            term_extractor = PhraseExtractor(
                term_phrases + abbreviations
            )
            for t in term_extractor.parse(tex_path, file_contents.contents):

                # Don't save term references if they are already in the definiendums.
                if any([overlaps(definiendum, t) for definiendum in all_definiendums]):
                    continue

                logging.debug(
                    "Found reference to term %s at (%d, %d) in %s for arXiv ID %s",
                    t.text,
                    t.start,
                    t.end,
                    t.tex_path,
                    item.arxiv_id,
                )
                type_ = (
                    "abbreviation"
                    if t.text in abbreviations
                    else "term"
                    if t.text in term_phrases
                    else "symbol"
                    if t.text in symbol_nicks
                    else "unknown"
                )
                yield TermReference(
                    id_=f"term-{t.tex_path}-{term_index}",
                    text=t.text,
                    type_=type_,
                    definition_ids=definition_ids[t.text],
                    definitions=definition_texts[t.text],
                    definition_texs=definition_texs[t.text],
                    sources=sources[t.text],
                    position_ratios=position_ratios[t.text],
                    section_names=section_names[t.text],
                    start=t.start,
                    end=t.end,
                    tex_path=t.tex_path,
                    tex=t.tex,
                    context_tex=t.context_tex,
                )
                term_index += 1

    def save(
        self,
        item: DetectDefinitionsTask,
        result: Union[Definiendum, Definition, TermReference],
    ) -> None:

        output_dir = directories.arxiv_subdir("detected-definitions", item.arxiv_id)
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        definiendums_path = os.path.join(output_dir, "entities-definiendums.csv")
        definitions_path = os.path.join(output_dir, "entities-definitions.csv")
        term_references_path = os.path.join(output_dir, "entities-term-references.csv")

        if isinstance(result, Definiendum):
            file_utils.append_to_csv(definiendums_path, result)
        elif isinstance(result, Definition):
            file_utils.append_to_csv(definitions_path, result)
        elif isinstance(result, TermReference):
            file_utils.append_to_csv(term_references_path, result)


import logging
import numpy as np
import os.path
import pandas as pd
import pyrsktools

from burst import BurstFourier, BurstWelch, BurstPeaks
from constants import (H5_PATH, OUTPUT_PATH, PROCESSED_PATH, VARIABLES,
                       TIMEZONE, AVG_FOLDER, Z_ELEVATION, DEVICES)
from intervals import DATA_INTERVALS, CALM_INTERVALS, STORM_INTERVALS
from tools import plotter, station

class Device(object):
    r"""
    Represents a RbR Concerto and its associated data (if loaded)

    Parameters
    ----------
    site : str
        ['S1', .. 'S5']
    dtype : string
        'floater' or 'bedframe'
    file : str
        file name without extension
    f : int
        Sampling frequency [Hz]
    sr : int
        Sampling rate (samples per burst)
    i : int
        Interval between bursts [s]
    T : array_like
        Turbidity calibration values [NTU]
    SSC : array_like
        Suspended sediment concentration [mg/L]
        Same length as T
    dformat : string
        Data format "h5" or "rsk"
    """

    def __init__(self, site, dtype, file, f, sr, i, T, SSC, dformat):
        self.site = site
        self.dtype = dtype
        self._init_logger()
        if dtype not in ["floater", "bedframe"]:
            raise ValueError("Unknown device type")
        if dformat not in ["rsk", "h5"]:
            raise ValueError("Unknown data format type")
        self.file = file
        self.f = f
        self.sr = sr
        self.i = i
        self.z = Z_ELEVATION
        self.df_tidal = None
        if len(T) != len(SSC):
            raise ValueError("T and SSC must have the same length")
        self.T = T
        self.SSC = SSC
        self.format = dformat
        self.vars = []
        self._load_data()

    def _init_logger(self):
        self.logger = logging.getLogger(str(self))
        self.logger.setLevel(logging.INFO)
        if not self.logger.handlers:
            fh = logging.FileHandler("./logs/%s.log" % str(self))
            fh.setLevel(logging.INFO)
            formatter = logging.Formatter(
                "%(asctime)s - %(name)s - %(levelname)s - %(message)s")
            fh.setFormatter(formatter)
            self.logger.addHandler(fh)
        self.logger.info("Device %s create" % str(self))

    def _set_vars(self):
        """
        Set number of variables available in self.df pandas.DataFrame
        """
        if self.vars == []:
            for v in self.df.columns:
                if v in VARIABLES.keys():
                    self.vars.append(VARIABLES[v])

    def _load_data(self):
        """
        Loads data from device file into self.df as a pandas.DataFrame.
        Loads burst-averaged data into self.df_avg as a pandas.DataFrame.
        """
        self.df = None
        self.df_avg = None
        if self.format == "h5":
            data_path = self.get_H5_path()
            self.df = pd.read_hdf(data_path, "df")
            try:
                self.df_avg = pd.read_hdf(self.get_H5_avg_path(), "df")
            except FileNotFoundError:
                self.set_df_avg(save=True)
            self.set_tide()
            self.logger.info(
                "Using %s as a dataframe source for averaged %s",
                self.get_H5_avg_path(),
                str(self))
            self.logger.info(self.df_avg.shape)
            self.logger.info(
                "NAN values: %s",
                str(self.df_avg.isnull().T.any().T.sum()))
        else:  # rsk
            data_path = self.get_RSK_path()
            rsk = pyrsktools.open(data_path)
            # Pandas dataframe for the win
            self.df = pd.DataFrame(rsk.npsamples())
            # timestamp as index, localize it as NZST
            self.df = self.df.set_index("timestamp")
            self.df.index = self.df.index.tz_convert(None)
            self.df.index = self.df.index.tz_localize(TIMEZONE)
            self.set_ssc()
        self.logger.info(
            "Using %s as a dataframe source for %s",
            data_path,
            str(self))
        self.logger.info(self.df.shape)
        self.logger.info(
            "NAN values: %s",
            str(self.df.isnull().T.any().T.sum()))
        self._set_vars()

    def _calc_bursts(self):
        """
        Calculates U, T and H for each valid burst.
        Note: only available for bedframe devices.
        """
        # Orbital speed, sig wave height and period
        if self.dtype != "bedframe":
            raise NotImplementedError("Not available in current class")
        self.df_avg["u"] = np.NaN
        self.df_avg["T"] = np.NaN
        self.df_avg["H"] = np.NaN
        df = self.clean_df(self.df, average=False)
        # Calculate U for each available burst
        start_date = self.df_avg.index[0]
        while start_date <= self.df_avg.index[-1]:
            end_date = start_date + pd.Timedelta("%ss" % self.i)
            burst = self.get_burst(
                        start=start_date,
                        end=end_date,
                        method="welch",
                        df=df)
            if burst:
                self.df_avg.loc[start_date, ["u", "T", "H"]] = burst.get_UTH()
            else:
                self.logger.warning("Invalid burst start date %s end date %s",
                                    str(start_date),
                                    str(end_date))
            start_date = end_date

    def set_tide(self):
        """
        Calculate trend of tide Ebb/Flood depending on next row's depth
        """
        if "Tide" not in self.df_avg.columns:
            self.df_avg["Tide"] = np.where(
                self.df_avg["depth_00"] > self.df_avg["depth_00"].shift(-1),
                "Ebb",
                "Flood")
        # last row doesn't have a next, compare with previous row
        last_rows = self.df_avg.tail(2)
        if last_rows.iloc[0].depth_00 > last_rows.iloc[1].depth_00:
            self.df_avg.at[self.df_avg.index[-1], "Tide"] = "Ebb"
        else:
            self.df_avg.at[self.df_avg.index[-1], "Tide"] = "Flood"

    def set_ssc(self):
        """
        Calculate SSC from linear interpolation of self.T with self.SSC
        """
        if "ssc" not in self.df.columns:
            self.df["ssc"] = self.df.apply(
                lambda r: np.interp(
                    r.turbidity_00,
                    self.T,
                    self.SSC),
                axis=1)
            # S1 floater saturated, remove max values
            if self.site == "S1" and self.dtype == "floater":
                d = self.get_dict_device()
                self.df.loc[
                    self.df["ssc"] == d["ssc_saturated_value"], ["ssc"]
                    ] = np.NaN

    def save_H5(self, avg=False):
        """
        Saves device data to h5 file
        """
        if not os.path.isfile(self.get_H5_path()):
            self.df.to_hdf(self.get_H5_path(), key="df", mode="w")
        if avg:
            self.df_avg.to_hdf(self.get_H5_avg_path(), key="df", mode="w")

    def get_H5_path(self):
        return "%s%s.h5" % (H5_PATH, self.file)

    def get_H5_avg_path(self):
        return "%s%s/%s.h5" % (H5_PATH, AVG_FOLDER, self.file)

    def get_RSK_path(self):
        return "%s%s_processed.rsk" % (PROCESSED_PATH, self.file)

    def get_burst(self, start=None, end=None, method="fourier", df=None):
        """"
        Get Burst from start to end dates

        """
        if df is None:
            df = self.df
        bvars = [
            "salinity_00",
            "temperature_00",
            "seapressure_00",
            "depth_00"]
        dfburst = df[(df.index >= start) & (df.index < end)][:self.sr][bvars]
        # discard burst with missing values or NaN
        if ((len(dfburst) < (self.sr/2)) or
            (dfburst.isnull().values.sum() != 0) or
                (dfburst.isna().values.sum() != 0)):
            return None
        if method == "fourier":
            return BurstFourier(
                dfburst, dfburst.index, self.f, self.z, str(self))
        elif method == "welch":
            return BurstWelch(
                dfburst, dfburst.index, self.f, self.z, str(self))
        else:
            return BurstPeaks(
                dfburst, dfburst.index, self.f, self.z, str(self))

    def set_df_avg(self, save=False):
        """
        Calculates SSC, clean data, average and save pandas.DataFrame
        in self.data_path_avg file
        """
        self.set_ssc()
        self.df_avg = self.clean_df(self.df)
        self.df_avg["ssc_sd"] = self.df.ssc.resample("%ss" % self.i).std()
        if self.dtype == "bedframe":
            self._calc_bursts()
        self.save_H5(avg=save)

    def clean_df(self, df, average=True):
        """ Clean and average given dataframe df """
        intervals = DATA_INTERVALS[self.__str__()]
        if average:
            df = df.resample("%ss" % self.i, label="left").mean()
            dfr = pd.DataFrame(
                    index=df.index,
                    columns=df.columns) if intervals else df
        else:
            dfr = pd.DataFrame(columns=df.columns) if intervals else df
        if intervals:
            for interval in intervals:
                if average:
                    dfr.update(df[interval[0]:interval[1]])
                else:
                    dfr = dfr.append(df[interval[0]:interval[1]])
        # just in case some dodgy data sneaked into the intervals
        dfr.loc[dfr.turbidity_00 < 0, "turbidity_00"] = np.nan
        return dfr

    def get_depth_stats(self):
        """
        Burst-averaged mean, max, min values of Depth [m]
        """
        return (round(self.df_avg.depth_00.mean(), 2),
                round(self.df_avg.depth_00.max(), 2),
                round(self.df_avg.depth_00.min(), 2))

    def get_wave_stats(self):
        """
        Burst-averaged mean, max, min values of sig wave height [m]
        """
        return (round(self.df_avg[self.df_avg.H > 0].H.mean(), 2),
                round(self.df_avg[self.df_avg.H > 0].H.max(), 2),
                round(self.df_avg[self.df_avg.H > 0].H.min(), 2))

    def get_period_stats(self):
        """
        Burst-averaged mean, max, min values of peak period [s]
        """
        return (round(self.df_avg[self.df_avg["T"] > 0]["T"].mean(), 2),
                round(self.df_avg[self.df_avg["T"] > 0]["T"].max(), 2),
                round(self.df_avg[self.df_avg["T"] > 0]["T"].min(), 2))

    def get_u_stats(self):
        """
        Burst-averaged mean, max, min values of orbital velocity [cm/s]
        """
        return (round(self.df_avg[self.df_avg["u"] > 0]["u"].mean(), 2),
                round(self.df_avg[self.df_avg["u"] > 0]["u"].max(), 2),
                round(self.df_avg[self.df_avg["u"] > 0]["u"].min(), 2))

    def get_ssc_stats(self):
        """
        Burst-averaged mean, max, min values of SSC [mg/l]
        """
        return (round(self.df_avg[self.df_avg["ssc"] > 0.01]["ssc"].mean(), 4),
                round(self.df_avg[self.df_avg["ssc"] > 0.01]["ssc"].max(), 4),
                round(self.df_avg[self.df_avg["ssc"] > 0.01]["ssc"].min(), 4))

    def get_time_stats(self):
        """
        Get % of time in the water (depth available)
        """
        perc = (len(self.df_avg[self.df_avg.depth_00 > 0.025]) /
                len(self.df_avg)) * 100
        self.logger.info("Time in the water %f", perc)
        return round(perc, 2)

    def get_dict_device(self):
        """ Return dict element defined in constants file """
        return next(item for item in DEVICES if (item["site"] == self.site and
                    item["type"] == self.dtype))

    def __str__(self):
        return ("%s %s") % (self.site, self.dtype.capitalize())

    def unicode(self):
        return self.__str__()

    def plot_days(self):
        dflist = [group for group in self.df.groupby(self.df.index.date)]
        for date, dfday in dflist:
            # all variables in the same plot
            dest_file = "%s%s/%s/%s.png" % (
                OUTPUT_PATH,
                self.site,
                self.dtype,
                str(date))
            plotter.plot_all_hourly(dest_file, date, dfday, self.vars)
            title = "%s %s" % (str(self), str(date))
            # averaged turb and depth (W/O cleaning it)
            dest_file = "%s%s/%s/%s/%s.png" % (
                OUTPUT_PATH,
                self.site,
                self.dtype,
                AVG_FOLDER,
                str(date))
            dfr = dfday.resample("%ss" % self.i).mean()
            plotter.plot_hourly_ssc_depth_avg(
                dfr[["ssc", "depth_00"]],
                date,
                dest_file,
                title)
            if self.df_avg is not None:
                # averaged CLEAN turb and depth
                dfr = self.df_avg[self.df_avg.index.date == date]
                parts = dest_file.split(".")
                dest_file = ".".join(parts[:-1]) + "_clean" + "." + parts[-1]
                plotter.plot_hourly_ssc_depth_avg(
                    dfr[["ssc", "depth_00"]],
                    date,
                    dest_file,
                    title)

    def plot_avg(self):
        """
        Plot averaged-bursts of SSC vs depth and salinity
        """
        dest_file = "%s%s/%s/%s/ssc.png" % (
                OUTPUT_PATH,
                self.site,
                self.dtype,
                AVG_FOLDER)
        plotter.plot_ssc_avg(self.df_avg, dest_file, str(self))

    def plot_ssc_u(self):
        """ Plots SSC vs wave orbital velocity """
        dest_file = "%s%s/%s/%s/%s_%s_ssc_vs_u.png" % (
            OUTPUT_PATH,
            self.site,
            self.dtype,
            AVG_FOLDER,
            self.site.lower(),
            self.dtype.lower())
        plotter.plot_ssc_u(self.df_avg, dest_file, str(self))
        dest_file = "%s%s/%s/%s/%s_%s_ssc_vs_u_log.png" % (
            OUTPUT_PATH,
            self.site,
            self.dtype,
            AVG_FOLDER,
            self.site.lower(),
            self.dtype.lower())
        plotter.plot_ssc_u_log(self.df_avg, dest_file, str(self))

    def plot_ssc_u_h(self, dffl=None):
        """ Plots a time series for SSC, Sig wave height and water depth """
        dest_file = "%s%s/%s/%s/%s_ssc_u_h_series.png" % (
            OUTPUT_PATH,
            self.site,
            self.dtype,
            AVG_FOLDER,
            str(self).lower())
        plotter.plot_ssc_u_h_series(self.df_avg, dffl, dest_file, str(self))

    def plot_ssc_u_h_weekly(self, dffl=None):
        """ Plots a detailed weekly times series """
        dfwindlist = station.get_weekly_wind()
        dfrainlist = station.get_weekly_rainfall()
        dfpressurelist = station.get_weekly_pressure()
        dfrivers = station.get_weekly_rivers()
        dflist = [g for g in self.df_avg.groupby(self.df_avg.index.week)]
        if dffl is not None:
            dffllist = [g for g in dffl.groupby(dffl.index.week)]
        else:
            dffllist = [[None, None] for i in dflist]
        depthlist = self.get_weekly_corrected_depth()
        i = 0
        for date, dfweek in dflist:
            dfriverlist = [(k, val[i]) for k, val in dfrivers.items()]
            dest_file = "%s%s/%s/%s/weekly_ssc_u_h_series_%d.png" % (
                OUTPUT_PATH,
                self.site,
                self.dtype,
                AVG_FOLDER,
                i)
            dfweek["u"] = dfweek["u"].fillna(-1)
            plotter.plot_ssc_u_h_weekly_series(
                dfweek,
                dffllist[i][1],
                dfwindlist[i],
                dfrainlist[i],
                dfpressurelist[i],
                depthlist[i],
                dfriverlist,
                dest_file,
                date,
                str(self),
                i)
            i += 1

    def plot_depth(self):
        return plotter.plot_depth(self.df_avg)

    def get_weekly_corrected_depth(self):
        self.df["depth_corrected"] = self.df["depth_00"] - 0.15
        self.df.loc[self.df["depth_corrected"] < 0, ["depth_corrected"]] = 0
        return [g[1]["depth_corrected"].resample("%ss" % self.i).mean()
                for g in self.df.groupby(self.df.index.week)]

    def plot_tidal_ssc(self, intervals=None):
        """
        Plot SSC and U values per tidal cycle defined by given intervals.
        """
        if intervals is None:
            intervals = STORM_INTERVALS[self.__str__()]
        else:
            intervals = intervals[self.__str__()]
        tidal_vars = ["hours", "u", "ssc", "H", "depth_00", "T"]

        if self.df_tidal is None:
            tdelta = pd.Timedelta("300s")
            df_tidal = pd.DataFrame()
            for interval in intervals:
                df_int = self.df_avg[interval[0]:interval[1]]
                mid = df_int[df_int["Tide"] == "Ebb"].head(1).index[0] - tdelta
                df_int["hours"] = df_int.apply(
                    lambda r: get_hours(r, mid),
                    axis=1)
                df_tidal = df_tidal.append(df_int[tidal_vars])
                df_tidal["ssc"] = pd.to_numeric(df_tidal["ssc"])
            self.df_tidal = df_tidal
        plotter.plot_tidal_u_ssc(self.df_tidal, intervals)

    def get_intervals_df(self, intervals=None):
        """
        Plot SSC and U values per tidal cycle defined by given intervals.
        """
        if intervals is None:
            intervals = STORM_INTERVALS[self.__str__()]
        else:
            intervals = intervals[self.__str__()]
        tidal_vars = ["u", "ssc", "depth_00", "Tide"]

        if self.df_tidal is None:
            tdelta = pd.Timedelta("300s")
            df_tidal = pd.DataFrame()
            for interval in intervals:
                df_int = self.df_avg[interval[0]:interval[1]]
                # mid = df_int[df_int["Tide"] == "Ebb"].head(1).index[0] - tdelta
                # df_int["hours"] = df_int.apply(
                #     lambda r: get_hours(r, mid),
                #     axis=1)
                df_tidal = df_tidal.append(df_int[tidal_vars])
                df_tidal["ssc"] = pd.to_numeric(df_tidal["ssc"])
                df_tidal["depth_00"] = pd.to_numeric(df_tidal["depth_00"])
                df_tidal["site"] = self.site
                self.df_tidal = df_tidal
        return self.df_tidal

def get_hours(row, mid_value):
    """
    Get signed difference, in seconds, between row index (name)
    and given mid_value timestamp
    """
    if (row.name < mid_value):
        return -(mid_value - row.name).total_seconds()/3600
    else:
        return (row.name - mid_value).total_seconds()/3600


import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.metrics import accuracy_score, f1_score, classification_report, balanced_accuracy_score
import xgboost as xgb
import pickle

reviews = pd.read_csv('reviews.csv')
labels = pd.read_csv('labels.csv')

X = reviews['text']
y = labels['sentiment']

# Converting the Positive/Negative to 0/1 labels
mapping = {'Positive': 1, 'Negative': 0}
y = y.map(mapping)

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42, stratify=y)

vectorizer = CountVectorizer()
X_train_vec = vectorizer.fit_transform(X_train)
X_test_vec = vectorizer.transform(X_test)

with open('vectorizer.pkl', 'wb') as model_file:
    pickle.dump(vectorizer, model_file)

xgb_classifier = xgb.XGBClassifier()
xgb_classifier.fit(X_train_vec, y_train)

with open('xgb_classifier.pkl', 'wb') as model_file:
    pickle.dump(xgb_classifier, model_file)

y_pred = xgb_classifier.predict(X_test_vec)

# Calculating balanced accuracy
balanced_acc = balanced_accuracy_score(y_test, y_pred)

# Printing the result
print("Balanced Accuracy:", balanced_acc)

# Additional metrics
accuracy = accuracy_score(y_test, y_pred)
macro_f1 = f1_score(y_test, y_pred, average='macro')
report = classification_report(y_test, y_pred)

print(f"Accuracy: {accuracy}")
print(f"Macro-average F1 Score: {macro_f1}")
print("Classification Report:\n", report)


import lxml.html as le
from lxml import etree
import re
import requests
import os
from time import sleep
import csv
from datetime import datetime

def RemoveNonAscii(text):
    res=''.join([i if ord(i) < 128 else ' ' for i in text]).strip()
    if(res==''):
        return '0'
    return res

nameRe="\"nom\":\"(.*?)\","
priceRe="\"prixPromo\":\"(.*?)\""

headerName='-----NAME-------'
headerPrice='--------PRICE---------'

url_dict={
    "site_1_2_Url":"http://mobile.free.fr/",
    "site_1_1_Url":"https://www.free.fr/freebox/",
    "site_2_1_Url":"https://www.sfr.fr/offre-internet/box/#sfrintid=HS_MAIN_Offres-box_PAR",
    "site_2_2_Url":"https://www.sfr.fr/forfait-mobile/offres/forfait-mobile#sfrintid=HS_MAIN_Forfaits-mobile_PAR",
    "site_3_1_Url":"https://www.bouyguestelecom.fr/offres-internet",
    "site_3_2_Url":"https://www.bouyguestelecom.fr/bons-plans/serie-speciale-b-you",
    "site_4_1_Url":"https://boutique.orange.fr/internet/offres-fibre",
    "site_4_2_Url":"https://boutique.orange.fr/mobile/forfaits-orange"
}

xPathDict={"name":"//h5/div/div[@class='markdown-p']",\
    "w-div":"//div[@class='sc-hmXxxW fuYjca']/span[@class='sc-jTzLTM iXZCuB']",
    "b-div":"//div[@class='sc-hmXxxW jPKUCL']/span[@class='sc-jTzLTM iXZCuB']",
    "test":"//div[@class='sc-hmXxxW jPKUCL']/span[@class='sc-fjdhpX jZqwup']"
    }
currDir = os.path.dirname(os.path.realpath(__file__))
dataList=[]

def scrapeSite1():
    #part1
    print("1-1 Started")
    s=requests.session()
    html=s.get(url_dict['site_1_1_Url']).text
    tree=le.fromstring(html)
    names=tree.xpath(xPathDict['name'])
    wd=tree.xpath(xPathDict['w-div'])
    bd=tree.xpath(xPathDict['b-div'])
    prices_div=bd+wd
    prices=[]
    for p in prices_div:
        price=""
        left=p.findall("./span[@class='sc-kAzzGY hqcLUv']")
        right=p.findall("./span[@class='sc-fjdhpX jZqwup']/span[@class='sc-jzJRlG eLwrik']/span[@class='sc-kgoBCf kVSTZj']")
        if(len(left)==0):
            price="0.0"
        elif(len(left)!=0 and len(right)==0):
            price=left[0].text+".0"
        else:
            price=left[0].text+"."+right[0].text
        prices.append(price)
    if(len(names)==len(prices)):
        l=len(names)
        for i in range(l):
            data=dict()
            data[headerName]=names[i].text
            data[headerPrice]=prices[i]
            dataList.append(data)
    print("1-1 Completed")
    #Part 2 of site 1
    print("1-2 Started")
    s=requests.session()
    html=s.get(url_dict['site_1_2_Url']).text
    tree=le.fromstring(html)
    dataDivs=tree.xpath("//div[@class='grid-c cell-top forfait']")
    for div in dataDivs:
        name=div.xpath("./h1[@class='description']")
        name2=div.xpath("./h1[@class='description']/span[1]")
        left=div.xpath("./div[contains(@class,'prix-mensuel')]/span[@class='prix prix-red prix-main']")
        right=div.xpath("./div[contains(@class,'prix-mensuel')]/span[@class='prix-other red']/span[@class='prix prix-red prix-cent']")
        price=RemoveNonAscii(left[0].text.strip())+"."+RemoveNonAscii(right[0].text)
        #print(name[0].text+" "+name2[0].text+" "+price)
        data=dict()
        data[headerName]=str(name[0].text+" "+name2[0].text).strip()
        data[headerPrice]=price
        dataList.append(data)
    print("1-2 Completed")

def scrapeSite2():
    print("2-1 Started")
    s=requests.session()
    html=s.get(url_dict['site_2_1_Url']).text
    tree=le.fromstring(html)
    divs=tree.xpath("//div[@class='overviewCol__inner']")
    for div in divs:
        name=div.findall(".//h2[@class='title600']")
        left=div.findall(".//div[@class='prix__valeur']")
        right=div.findall(".//span[@class='prix__devise']")
        price=RemoveNonAscii(left[0].text.strip())+"."+RemoveNonAscii(right[0].text)
        #print(name[0].text+" "+price)
        data=dict()
        data[headerName]=name[0].text
        data[headerPrice]=price
        dataList.append(data)
    print("2-1 Completed")

    #part 2
    print("2-2 Started")
    s=requests.session()
    html=s.get(url_dict['site_2_2_Url']).text
    tree=le.fromstring(html)
    divs=tree.xpath("//div[@class='overviewCol__inner']")
    for div in divs:
        name=div.findall(".//h2[@class='title600']")
        left=div.findall(".//div[@class='prix__valeur']")
        right=div.findall(".//span[@class='prix__devise']")
        price=RemoveNonAscii(left[0].text.strip())+"."+RemoveNonAscii(right[0].text)
        #print("2-2-Mob-  "+name[0].text+" "+price)
        data=dict()
        data[headerName]=name[0].text
        data[headerPrice]=price
        dataList.append(data)
    print("2-2 Completed")

def scrapeSite3():
    #PART 1
    print("3-1 Started")
    s=requests.session()
    html=s.get(url_dict['site_3_1_Url']).text
    tree=le.fromstring(html)
    divs=tree.xpath("//div[contains(@class,'column is-half-tablet is-3-desktop')]")
    for div in divs:
        name=div.findall(".//title")
        if(len(name)==0):
            img=div.findall(".//div/header/h2/img")
            name=str(img[0].get('alt'))
        left=div.findall(".//p[@class='main']")
        right=div.findall(".//span[@class='centimes']")
        try:
         price=RemoveNonAscii(left[0].text.strip())+"."+RemoveNonAscii(right[0].text)
        except:
         price="0.0"
        if(isinstance(name,str)):
          #print(name+" "+price)
          data=dict()
          data[headerName]=name
          data[headerPrice]=price
          dataList.append(data)
        else:
           #print(name[0].text+" "+price)
           data=dict()
           data[headerName]=name[0].text
           data[headerPrice]=price
           dataList.append(data)
    print("3-1 Completed")
    #PART 2
    print("3-2 Started")
    s=requests.session()
    html=s.get(url_dict['site_3_2_Url']).text
    tree=le.fromstring(html)
    divs=tree.xpath("//div[contains(@class,'offers')]/div[contains(@class,'offer')]")
    for div in divs:
        name=div.xpath(".//div[@class='head']/div[@class='head_container']/h2[contains(@class,'head_title')]")
        name=str(le.tostring(name[0],method="text").strip())[4:-1]
        left=div.findall(".//th[@class='big']")
        right=div.findall(".//th[@class='sup']")
        price=str(left[0].text)+"."+RemoveNonAscii(str(right[0].text_content()))
        data=dict()
        data[headerName]=name
        data[headerPrice]=price
        dataList.append(data)
    print("3-2 Completed")

def scrapeSite4():
    print("4-1 Started")
    s=requests.session()
    html=s.get(url_dict['site_4_1_Url']).text
    tree=le.fromstring(html)
    divs=tree.xpath("//div[@class='template header-crem']")
    for div in divs:
        name=div.findall(".//h3/span")
        left=div.findall(".//span[@class='integer']")
        right=div.findall(".//div[@class='price text-primary']/sup")
        price=RemoveNonAscii(left[0].text.strip())+"."+RemoveNonAscii(right[0].text).replace(",",'')
       # print(name[0].text+"   "+price)
        data=dict()
        data[headerName]=name[0].text
        data[headerPrice]=price
        dataList.append(data)
    print("4-1 Completed")
    #PART 2
    print("4-2 Started")
    s=requests.session()
    html=s.get(url_dict['site_4_2_Url']).text
    names=re.findall(nameRe,html)
    prices=re.findall(priceRe,html)
    if(len(names)==len(prices)):
        for i in range(len(names)):
            #print(names[i]+"--->"+prices[i])
            data=dict()
            data[headerName]=names[i]
            data[headerPrice]=prices[i].replace(',','.')
            dataList.append(data)
    print("4-2 Completed")

def init():
    try:
        dat=datetime.date(datetime.now())
        di={headerName:dat,headerPrice:"----"}
        dataList.append(di)
        di=currDir+"/data"
        os.mkdir(di)
    except:
        print("Folder already exist")

def writeListToCsv(toCSV):
    oldData=[]
    keys = toCSV[0].keys()
    fileName=currDir+"/data/data.csv"
    xists = os.path.isfile(fileName)
    if(xists):
        with open(fileName, "r") as f:
            reader = csv.DictReader(f)
            oldData=list(reader)
    #appending old data with new Data
    toCSV=toCSV+oldData
    with open(fileName, 'w') as output_file:
        dict_writer = csv.DictWriter(output_file, keys)
        dict_writer.writeheader()
        dict_writer.writerows(toCSV)

def main():
    init()
    scrapeSite1()
    scrapeSite2()
    scrapeSite3()
    scrapeSite4()
    writeListToCsv(dataList)

if __name__ == "__main__":
    main()


from ..utils import ext_loader

ext_module = ext_loader.load_ext('_ext', ['box_iou_rotated'])

def box_iou_rotated(bboxes1, bboxes2, mode='iou', aligned=False):
    """Return intersection-over-union (Jaccard index) of boxes.

    Both sets of boxes are expected to be in
    (x_center, y_center, width, height, angle) format.

    If ``aligned`` is ``False``, then calculate the ious between each bbox
    of bboxes1 and bboxes2, otherwise the ious between each aligned pair of
    bboxes1 and bboxes2.

    Arguments:
        boxes1 (Tensor): rotated bboxes 1. \
            It has shape (N, 5), indicating (x, y, w, h, theta) for each row.
            Note that theta is in radian.
        boxes2 (Tensor): rotated bboxes 2. \
            It has shape (M, 5), indicating (x, y, w, h, theta) for each row.
            Note that theta is in radian.
        mode (str): "iou" (intersection over union) or iof (intersection over
            foreground).

    Returns:
        ious(Tensor): shape (N, M) if aligned == False else shape (N,)
    """
    assert mode in ['iou', 'iof']
    mode_dict = {'iou': 0, 'iof': 1}
    mode_flag = mode_dict[mode]
    rows = bboxes1.size(0)
    cols = bboxes2.size(0)
    if aligned:
        ious = bboxes1.new_zeros(rows)
    else:
        ious = bboxes1.new_zeros((rows * cols))
    bboxes1 = bboxes1.contiguous()
    bboxes2 = bboxes2.contiguous()
    ext_module.box_iou_rotated(
        bboxes1, bboxes2, ious, mode_flag=mode_flag, aligned=aligned)
    if not aligned:
        ious = ious.view(rows, cols)
    return ious


import numpy as np
import os
import platform
from sklearn.metrics import accuracy_score, confusion_matrix
import matplotlib.pyplot as pl
import pandas as pd

from configs import *
from params import *
from generator import DataGenerator
from model import loss_asc, loss_dc, create_model
from normalization import mezza_normalization

""" Script to run inference of all trained models on the test data """

class_labels_asc = ['airport', 'shopping_mall', 'metro_station', 'street_pedestrian',
                    'public_square', 'street_traffic', 'tram', 'bus', 'metro', 'park']

class_labels_dc = ['A', 'B', 'C']

class ASCPredictor:
    """ Prediction class """
    def __init__(self,
                 fn_model,
                 do_masking=False,
                 batch_size=512,
                 emb_dim=256):
        """ Initialize class
        Args:
            fn_model (str): Model filename (note that we only store the weights!, saving the full model
                            instead caused some errors with the custom loss functions)
            do_masking (bool): Switch to do embedding masking
            batch_size (int): Batch size
            emb_dim (int): Embedding size
        """
        self.emb_dim = emb_dim
        self.do_masking = do_masking

        shape_in = (batch_size, 64, 431, 1)

        # create model from scratch and load weights
        self.multi_branch_model = create_model(shape_in=shape_in,
                                               n_output_dim=10,
                                               do_masking=do_masking,
                                               emb_dim=self.emb_dim,
                                               inference=True)
        self.multi_branch_model.load_weights(fn_model)

        # double embedding size if embedding masking is used
        if do_masking:
            self.emb_dim *= 2

    def predict_feat(self, feat, do_activate_upper_embedding_half=True):
        """ Compute model prediction for given feature tensor
        Args:
            feat (4d np.ndarray): Feature tensor (batch x freq x time x channel)
            do_activate_upper_embedding_half (bool): If True and if masking is activated, binary mask is set such that
                                    embedding is optimized for ASC task
        Returns:
            prediction (2d np.ndarray): Prediction (batch x n_classes)
        """
        n_patches = feat.shape[0]
        mask = np.ones((n_patches, self.emb_dim))
        if self.do_masking:
            # if masking is selected, only consider first half of embedding (trained for ASC task)
            if do_activate_upper_embedding_half:
                mask[:, self.emb_dim//2:] = 0
            else:
                mask[:, :self.emb_dim//2] = 0

        # get output of dual-output model
        dual_predictions = self.multi_branch_model.predict((feat, mask))

        return dual_predictions

def create_cm_plot(fn_img, true, pred, labels):
    """ Create confusion matrix plot
    Args:
        fn_img (str): Image filename
        true (np.ndarray): True class indices
        pred (np.ndarray): Predicted class indices
    """
    cm = confusion_matrix(true, pred, normalize="true")
    tick_vals = np.arange(len(labels))
    pl.figure()
    pl.imshow(cm, aspect="equal")
    pl.xticks(tick_vals, labels, rotation=60)
    pl.yticks(tick_vals, labels)
    pl.tight_layout()
    pl.savefig(fn_img)
    pl.close()
    print(fn_img, ' saved')

if __name__ == '__main__':

    n_configs = len(configs)

    print("Import data")
    feat = {}
    asc_target = {}
    dc_target = {}
    is_source_domain = {}
    is_target_domain = {}

    # iterate over training and test set and import both
    for ds_type in ('training', 'test'):
        # let's use the data generator to import the data (but for nothing else)
        dg = DataGenerator(ds_type=ds_type,
                           dir_feat=dir_feat,
                           batch_size=batch_size)  # other parameters don't matter here

        # load data from generator
        feat[ds_type] = dg.feat
        asc_target[ds_type] = dg.asc_target
        dc_target[ds_type] = dg.dc_target
        is_source_domain[ds_type] = dc_target[ds_type][:, 0].astype(bool)
        is_target_domain[ds_type] = np.logical_not(is_source_domain[ds_type])

    print("Run predictions")
    for ds_type in ('training', 'test'):

        print('=' * 60)
        print(ds_type)
        print('=' * 60)

        # split into source (device A) and target (devices B & C)
        feat_source = feat[ds_type][is_source_domain[ds_type], :, :, :]
        feat_target = feat[ds_type][is_target_domain[ds_type], :, :, :]
        asc_true_s = asc_target[ds_type][is_source_domain[ds_type], :]
        asc_true_t = asc_target[ds_type][is_target_domain[ds_type], :]
        dc_true_s = dc_target[ds_type][is_source_domain[ds_type], :]
        dc_true_t = dc_target[ds_type][is_target_domain[ds_type], :]

        acc = np.zeros((n_configs, 8))

        # iterate over models and compute accuracy for source and target domain data in testset
        for c in range(n_configs):

            config_label = configs[c]["label"]

            print('-' * 60)
            print(c + 1, '/', n_configs, ':', config_label)
            print('-' * 60)

            fn_model = os.path.join(dir_results, 'model_{}.h5'.format(config_label))

            # if Mezza normalization is active, let's use the model trained with the same parameters
            fn_model = fn_model.replace('-mezza', '')

            # get prediction
            predictor = ASCPredictor(fn_model,
                                     do_masking=configs[c]["do_masking"],
                                     batch_size=batch_size)

            # Apply band-wise statistics matching procedure proposed by Mezza et al. (EUSIPCO 2021)
            if configs[c]["do_mezza"] and ds_type == "test":
                # get source domain data from training set
                feat_source_training = feat["training"][is_source_domain["training"], :, :, :]

                # adapt test data (target domain) to training data (source domain)
                feat_target[:, :, :, 0] = mezza_normalization(feat_target[:, :, :, 0],
                                                              feat_source_training[:, :, :, 0],
                                                              switch_time_freq=True)

            # get ASC predictions on source domain data
            asc_pred_s_u, dc_pred_s_u = predictor.predict_feat(feat_source, True)
            # get DC predictions on source domain data
            asc_pred_s_d, dc_pred_s_d = predictor.predict_feat(feat_source, False)
            # get ASC predictions on target domain data
            asc_pred_t_u, dc_pred_t_u = predictor.predict_feat(feat_target, True)
            # get DC predictions on target domain data
            asc_pred_t_d, dc_pred_t_d = predictor.predict_feat(feat_target, False)

            # accuracies
            acc[c, 0] = accuracy_score(np.argmax(asc_true_s, axis=1), np.argmax(asc_pred_s_u, axis=1))
            acc[c, 1] = accuracy_score(np.argmax(asc_true_s, axis=1), np.argmax(asc_pred_s_d, axis=1))
            acc[c, 2] = accuracy_score(np.argmax(asc_true_t, axis=1), np.argmax(asc_pred_t_u, axis=1))
            acc[c, 3] = accuracy_score(np.argmax(asc_true_t, axis=1), np.argmax(asc_pred_t_d, axis=1))

            acc[c, 4] = accuracy_score(np.argmax(dc_true_s, axis=1), np.argmax(dc_pred_s_u, axis=1))
            acc[c, 5] = accuracy_score(np.argmax(dc_true_s, axis=1), np.argmax(dc_pred_s_d, axis=1))
            acc[c, 6] = accuracy_score(np.argmax(dc_true_t, axis=1), np.argmax(dc_pred_t_u, axis=1))
            acc[c, 7] = accuracy_score(np.argmax(dc_true_t, axis=1), np.argmax(dc_pred_t_d, axis=1))

        # finally export CSV file with all metrics for all configs
        df_res = pd.DataFrame(data=acc, index=[_['label'] for _ in configs],
                              columns=('ASC-S-hi', 'ASC-S-lo', 'ASC-T-hi', 'ASC-T-lo', 'DC-S-hi',
                                       'DC-S-lo', 'DC-T-hi', 'DC-T-lo'))
        fn_csv = os.path.join(dir_results, 'acc_8col_{}.csv'.format(ds_type))
        df_res.to_csv(fn_csv, sep=';', float_format='%.2f')

    cmd = 'cp /mnt/IDMT-WORKSPACE/DATA-STORE/abr/asc_2021_domain_adaptation/results_audasc/* /home/avdata/xtract/abr/2021_icassp_asc_da/'
    print(cmd)
    os.system(cmd)

    print("Done :)")


import pygame, sys, random
# 初始坐标，确定行数和列数
block_initial_position, score, times, gameover, press, all_block, background = [20, 5], [0], 0, [], False, [
    [[0, 0], [0, -1], [0, 1], [0, 2]], [[0, 0], [0, 1], [-1, 1], [-1, 0]], [[0, 0], [0, -1], [-1, 0], [-1, 1]],
    [[0, 0], [0, 1], [-1, -1], [-1, 0]], [[0, 0], [0, 1], [1, 0], [0, -1]], [[0, 0], [1, 0], [-1, 0], [1, -1]],
    [[0, 0], [1, 0], [-1, 0], [1, 1]]], [[0 for column in range(0, 10)] for row in range(0, 22)]
background[0], select_block = [1 for column in range(0, 10)], list(random.choice(all_block))

# 移动函数处理
def move(n):
    # 设置一个初始值
    if n == 100:
        for row, column in select_block:
            pygame.draw.rect(screen, (255, 165, 0), (
                (column + block_initial_position[1]) * 40, 800 - (row + block_initial_position[0]) * 40, 38, 38))
        for row in range(0, 20):
            for column in range(0, 10):
                if background[row][column]: pygame.draw.rect(screen, (0, 0, 255), (column * 40, 800 - row * 40, 38, 38))
    y_drop, x_move = block_initial_position
    if n == 1 or n == -1:
        x_move += n
        for row, column in select_block:
            if (column + x_move) < 0 or (column + x_move) > 9 or background[column + y_drop][column + x_move]: break
        else:
            block_initial_position.clear(), block_initial_position.extend([y_drop, x_move])
    if n == 0:
        rotating_position = [(-column, row) for row, column in select_block]
        for row, column in rotating_position:
            if (column + x_move) < 0 or (column + x_move) > 9 or background[row + y_drop][row + x_move]: break
        else:
            select_block.clear(), select_block.extend(rotating_position)
    if n == 10:
        # 掉下的方块坐标-1
        y_drop -= 1
        for row, column in select_block:
            if background[row + y_drop][column + x_move] == 1: break
        else:
            block_initial_position.clear()
            block_initial_position.extend([y_drop, x_move])
            return
        for row, column in select_block: background[block_initial_position[0] + row][
            block_initial_position[1] + column] = 1
        complete_row = []
        for row in range(1, 20):
            if 0 not in background[row]: complete_row.append(x_move)
        complete_row.sort(reverse=True)
        for row in complete_row:
            background.pop(row)
            # 列表推导式
            background.append([0 for column in range(0, 10)])
        score[0] += len(complete_row)
        pygame.display.set_caption('俄罗斯方块[你的分数：' + str(score[0] * 100) + '分]')
        select_block.clear(), select_block.extend(list(random.choice(all_block)))
        block_initial_position.clear(), block_initial_position.extend([19, 4])
        for row, column in select_block:
            if background[row + block_initial_position[0]][column + block_initial_position[1]]: gameover.append(1)

# 初始化游戏
pygame.init()
screen = pygame.display.set_mode((400, 800))  # set_mode((400,800)) 修改数值可以修改窗口大小 但第7行和第10行要做相应的改动
while True:
    # 设置颜色
    screen.fill((255, 255, 255))
    # 按键处理
    for event in pygame.event.get():
        # 暂停游戏的处理
        if event.type == pygame.QUIT:
            sys.exit()
        # 方向键的处理 左下
        elif event.type == pygame.KEYDOWN and event.key == pygame.K_LEFT:
            move(-1)
        # 右下
        elif event.type == pygame.KEYDOWN and event.key == pygame.K_RIGHT:
            move(1)
        # 上下
        elif event.type == pygame.KEYDOWN and event.key == pygame.K_UP:
            move(0)
        elif event.type == pygame.KEYDOWN and event.key == pygame.K_DOWN:
            press = True
        elif event.type == pygame.KEYUP and event.key == pygame.K_DOWN:
            press = False
    if press: times += 10
    if times >= 50:
        # 移动10个单位长度
        move(10)
        times = 0
    else:
        # 时间处理
        times += 1
    # 游戏结束
    if gameover: sys.exit()
    move(100)
    pygame.time.Clock().tick(200)  # tick(200)修改数字可以修改游戏的整体速率
    pygame.display.flip()