#!/usr/bin/env python

from pyspark.sql import SparkSession

spark = SparkSession \
    .builder \
    .appName("Python Spark SQL basic example") \
    .config("spark.some.config.option", "some-value") \
    .getOrCreate()

# Load the Parquet file
on_time_dataframe = spark.read.parquet('data/on_time_performance.parquet')

# Use SQL to look at the total flights by month across 2015
on_time_dataframe.registerTempTable("on_time_dataframe")
total_flights_by_month = spark.sql(
"""Select Month, Year, Count(*) as total_flights
from on_time_dataframe
group by Year, Month
order by Year, Month"""
)

# This map/asDict trick makes the rows print a little prettier,
# It is optional
flights_chart_data = total_flights_by_month.rdd.map(lambda row: row.asDict())
print(flights_chart_data.collect())

# Save chart to MongoDB
import pymongo_spark
pymongo_spark.activate()
flights_chart_data.saveToMongoDB('mongodb://localhost:27017/agile_data_science.flights_by_month')

'''
[root@localhost workspace]# python total_flights.py
Setting default log level to "WARN".
To adjust logging level use sc.setLogLevel(newLevel). For SparkR, use setLogLevel(newLevel).
17/11/07 14:35:03 WARN NativeCodeLoader: Unable to load native-hadoop library for your platform... using builtin-java classes where applicable
17/11/07 14:35:03 WARN SparkConf:
SPARK_CLASSPATH was detected (set to '/root/workspace/lib/snappy-java-1.1.2.6.jar').
This is deprecated in Spark 1.0+.

Please instead use:
 - ./spark-submit with --driver-class-path to augment the driver classpath
 - spark.executor.extraClassPath to augment the executor classpath

17/11/07 14:35:03 WARN SparkConf: Setting 'spark.executor.extraClassPath' to '/root/workspace/lib/snappy-java-1.1.2.6.jar' as a work-around.
17/11/07 14:35:03 WARN SparkConf: Setting 'spark.driver.extraClassPath' to '/root/workspace/lib/snappy-java-1.1.2.6.jar' as a work-around.
17/11/07 14:35:03 WARN Utils: Your hostname, localhost.localdomain resolves to a loopback address: 127.0.0.1; using 192.168.1.103 instead (on interface eth0)
17/11/07 14:35:03 WARN Utils: Set SPARK_LOCAL_IP if you need to bind to another address
17/11/07 14:35:06 WARN Utils: Service 'SparkUI' could not bind on port 4040. Attempting port 4041.
SLF4J: Failed to load class "org.slf4j.impl.StaticLoggerBinder".
SLF4J: Defaulting to no-operation (NOP) logger implementation
SLF4J: See http://www.slf4j.org/codes.html#StaticLoggerBinder for further details.
[{'Year': '2015', 'Month': '1', 'total_flights': 469968}, {'Year': '2015', 'Month': '10', 'total_flights': 486165},
{'Year': '2015', 'Month': '11', 'total_flights': 467972}, {'Year': '2015', 'Month': '12', 'total_flights': 479230},
{'Year': '2015', 'Month': '2', 'total_flights': 429191}, {'Year': '2015', 'Month': '3', 'total_flights': 504312},
{'Year': '2015', 'Month': '4', 'total_flights': 485151}, {'Year': '2015', 'Month': '5', 'total_flights': 496993},
{'Year': '2015', 'Month': '6', 'total_flights': 503897}, {'Year': '2015', 'Month': '7', 'total_flights': 520718},
{'Year': '2015', 'Month': '8', 'total_flights': 510536}, {'Year': '2015', 'Month': '9', 'total_flights': 464946}]
'''

'''
[root@localhost workspace]# mongo agile_data_science
MongoDB shell version v3.4.2
connecting to: mongodb://127.0.0.1:27017/agile_data_science
MongoDB server version: 3.4.2
Server has startup warnings:
2017-08-09T15:01:54.581-0700 I STORAGE  [initandlisten]
2017-08-09T15:01:54.581-0700 I STORAGE  [initandlisten] ** WARNING: Using the XFS filesystem is strongly recommended with the WiredTiger storage engine
2017-08-09T15:01:54.581-0700 I STORAGE  [initandlisten] **          See http://dochub.mongodb.org/core/prodnotes-filesystem
2017-08-09T15:01:55.840-0700 I CONTROL  [initandlisten]
2017-08-09T15:01:55.840-0700 I CONTROL  [initandlisten] ** WARNING: Access control is not enabled for the database.
2017-08-09T15:01:55.840-0700 I CONTROL  [initandlisten] **          Read and write access to data and configuration is unrestricted.
2017-08-09T15:01:55.840-0700 I CONTROL  [initandlisten] ** WARNING: You are running this process as the root user, which is not recommended.
2017-08-09T15:01:55.840-0700 I CONTROL  [initandlisten]
2017-08-09T15:01:55.843-0700 I CONTROL  [initandlisten]
2017-08-09T15:01:55.843-0700 I CONTROL  [initandlisten] ** WARNING: /sys/kernel/mm/transparent_hugepage/enabled is 'always'.
2017-08-09T15:01:55.843-0700 I CONTROL  [initandlisten] **        We suggest setting it to 'never'
2017-08-09T15:01:55.844-0700 I CONTROL  [initandlisten]
2017-08-09T15:01:55.844-0700 I CONTROL  [initandlisten] ** WARNING: /sys/kernel/mm/transparent_hugepage/defrag is 'always'.
2017-08-09T15:01:55.844-0700 I CONTROL  [initandlisten] **        We suggest setting it to 'never'
2017-08-09T15:01:55.844-0700 I CONTROL  [initandlisten]
> db.flights_by_month.find().sort({"Year": 1, "Month": 1})
{ "_id" : ObjectId("599a09f14e9540530def1276"), "total_flights" : 469968, "Month" : "1", "Year" : "2015" }
{ "_id" : ObjectId("599a0aaf4e954054637f7559"), "total_flights" : 469968, "Month" : "1", "Year" : "2015" }
{ "_id" : ObjectId("5a0235464e954017f9eff004"), "total_flights" : 469968, "Month" : "1", "Year" : "2015" }
{ "_id" : ObjectId("599a09f34e9540530def1279"), "total_flights" : 486165, "Month" : "10", "Year" : "2015" }
{ "_id" : ObjectId("599a0ab04e954054637f755c"), "total_flights" : 486165, "Month" : "10", "Year" : "2015" }
{ "_id" : ObjectId("5a0235464e954017f9eff007"), "total_flights" : 486165, "Month" : "10", "Year" : "2015" }
{ "_id" : ObjectId("599a09f44e9540530def127c"), "total_flights" : 467972, "Month" : "11", "Year" : "2015" }
{ "_id" : ObjectId("599a0ab04e954054637f755f"), "total_flights" : 467972, "Month" : "11", "Year" : "2015" }
{ "_id" : ObjectId("5a0235464e954017f9eff00a"), "total_flights" : 467972, "Month" : "11", "Year" : "2015" }
{ "_id" : ObjectId("599a09f44e9540530def127f"), "total_flights" : 479230, "Month" : "12", "Year" : "2015" }
{ "_id" : ObjectId("599a0ab04e954054637f7562"), "total_flights" : 479230, "Month" : "12", "Year" : "2015" }
{ "_id" : ObjectId("5a0235464e954017f9eff00d"), "total_flights" : 479230, "Month" : "12", "Year" : "2015" }
{ "_id" : ObjectId("599a09f44e9540530def1282"), "total_flights" : 429191, "Month" : "2", "Year" : "2015" }
{ "_id" : ObjectId("599a0ab04e954054637f7565"), "total_flights" : 429191, "Month" : "2", "Year" : "2015" }
{ "_id" : ObjectId("5a0235464e954017f9eff010"), "total_flights" : 429191, "Month" : "2", "Year" : "2015" }
{ "_id" : ObjectId("599a09f44e9540530def1285"), "total_flights" : 504312, "Month" : "3", "Year" : "2015" }
{ "_id" : ObjectId("599a0ab04e954054637f7568"), "total_flights" : 504312, "Month" : "3", "Year" : "2015" }
{ "_id" : ObjectId("5a0235474e954017f9eff013"), "total_flights" : 504312, "Month" : "3", "Year" : "2015" }
{ "_id" : ObjectId("599a09f44e9540530def1288"), "total_flights" : 485151, "Month" : "4", "Year" : "2015" }
{ "_id" : ObjectId("599a0ab14e954054637f756b"), "total_flights" : 485151, "Month" : "4", "Year" : "2015" }
Type "it" for more
'''


"""
Describes the abstract Corpus class that all interfaces to corpora should
subclass.
"""

from collections import namedtuple
import logging.config
import os
from pathlib import Path
import pickle
import random
import subprocess
from typing import Any, List, Callable, Optional, Set, Sequence, Tuple, Type, TypeVar

import numpy as np

from .config import ENCODING
from .preprocess import feat_extract
from . import utils
from .exceptions import PersephoneException
from .exceptions import LabelMismatchException
from .preprocess import elan, wav
from . import utterance
from .utterance import Utterance
from .preprocess.labels import LabelSegmenter

logger = logging.getLogger(__name__) # type: ignore

CorpusT = TypeVar("CorpusT", bound="Corpus")

def ensure_no_set_overlap(train: Sequence[str], valid: Sequence[str], test: Sequence[str]) -> None:
    """ Ensures no test set data has creeped into the training set."""

    logger.debug("Ensuring that the training, validation and test data sets have no overlap")
    train_s = set(train)
    valid_s = set(valid)
    test_s = set(test)

    if train_s & valid_s:
        logger.warning("train and valid have overlapping items: {}".format(train_s & valid_s))
        raise PersephoneException("train and valid have overlapping items: {}".format(train_s & valid_s))
    if train_s & test_s:
        logger.warning("train and test have overlapping items: {}".format(train_s & test_s))
        raise PersephoneException("train and test have overlapping items: {}".format(train_s & test_s))
    if valid_s & test_s:
        logger.warning("valid and test have overlapping items: {}".format(valid_s & test_s))
        raise PersephoneException("valid and test have overlapping items: {}".format(valid_s & test_s))

def find_untranscribed_wavs(wav_path: Path, transcription_path: Path, label_type: str) -> List[str]:
    """Find the prefixes for all the wav files that do not have an associated transcription
    Args:
        wav_path: Path to search for wav files in
        transcription_path: Path to search for transcriptions in
        label_type: The type of labels for transcriptions. Eg "phonemes" "phonemes_and_tones"
    Returns:
        A list of all untranscribed prefixes
    """
    audio_files = wav_path.glob("**/*.wav")
    transcription_files = transcription_path.glob("**/*.{}".format(label_type))

    transcription_file_prefixes = [t_file.stem for t_file in transcription_files]

    untranscribed_prefixes = [] # type: List[str]
    for a_file in audio_files:
        if a_file.stem not in transcription_file_prefixes:
            untranscribed_prefixes.append(a_file.stem)
    return untranscribed_prefixes

def get_untranscribed_prefixes_from_file(target_directory: Path) -> List[str]:
    """
    The file "untranscribed_prefixes.txt" will specify prefixes which
    do not have an associated transcription file if placed in the target directory.

    This will fetch those prefixes from that file and will return an empty
    list if that file does not exist.

    See find_untranscribed_wavs function for finding untranscribed prefixes in an
    experiment directory.

    Returns:
        A list of all untranscribed prefixes as specified in the file
    """

    untranscribed_prefix_fn = target_directory / "untranscribed_prefixes.txt"
    if untranscribed_prefix_fn.exists():
        with untranscribed_prefix_fn.open() as f:
            prefixes = f.readlines()

        return [prefix.strip() for prefix in prefixes]
    else:
        #logger.warning("Attempting to get untranscribed prefixes but the file ({})"
        #                " that should specify these does not exist".format(untranscribed_prefix_fn))
        pass
    return []

class Corpus:
    """ Represents a preprocessed corpus that is ready to be used in model
    training.

    Construction of a `Corpus` instance involves preprocessing data if the data has
    not previously already been preprocessed. The extent of the preprocessing
    depends on which constructor is used. If the default constructor,
    `__init__()` is used, transcriptions are assumed to already be preprocessed
    and only speech feature extraction from WAV files is performed. In other
    constructors such as `from_elan()`, preprocessing of the transcriptions is
    performed. See the documentation of the relevant constructors for more
    information.

    Once a Corpus object is created it should be considered immutable. At this
    point feature extraction from WAVs will have been performed, with feature
    files in `tgt_dir/feat/`.  Transcriptions will have been segmented into
    appropriate tokens (labels) and will be stored in `tgt_dir/label/`.

    """

    def __init__(self, feat_type: str, label_type: str, tgt_dir: Path,
                 *,
                 labels: Optional[Set[str]] = None,
                 max_samples: int=1000,
                 speakers: Optional[Sequence[str]] = None) -> None:
        """ Construct a `Corpus` instance from preprocessed data.

        Assumes that the corpus data has been preprocessed and is
        structured as follows: (1) WAVs for each utterance are found in
        `<tgt_dir>/wav/` with the filename `<prefix>.wav`, where `prefix` is
        some string uniquely identifying the utterance; (2) For each WAV file,
        there is a corresponding transcription found in `<tgt_dir>/label/` with
        the filename `<prefix>.<label_type>`, where `label_type` is some string
        describing the type of label used (for example, "phonemes" or "tones").

        If the data is found in the format, WAV normalization and speech
        feature extraction will be performed during `Corpus` construction, and
        the utterances will be randomly divided into training, validation and
        test_sets. If you would like to define these datasets yourself, include
        files named `train_prefixes.txt`, `valid_prefixes.txt` and
        `test_prefixes.txt` in `<tgt_dir>`. Each file should be a list of
        prefixes (utterance IDs), one per line. If these are found during
        `Corpus` construction, those sets will be used instead.

        Args:
            feat_type: A string describing the input speech features. For
                       example, "fbank" for log Mel filterbank features.
            label_type: A string describing the transcription labels. For example,
                         "phonemes" or "tones".
            labels: A set of strings representing labels (tokens) used in
                    transcription. For example: {"a", "o", "th", ...}.
                    If this parameter is not provided the experiment directory is
                    scanned for labels present in the transcription files.
            max_samples: The maximum number of samples an utterance in the
                corpus may have. If an utterance is longer than this, it is not
                included in the corpus.

        """

        if speakers:
            raise NotImplementedError("Speakers not implemented")

        logger.debug("Creating a new Corpus object with feature type {}, label type {},"
                     "target directory {}, label set {}, max_samples {}, speakers {}".format(
                    feat_type, label_type, labels, tgt_dir, max_samples, speakers)
        )

        # In case path is supplied as a string, make it a Path
        self.tgt_dir = Path(tgt_dir)

        #: A string representing the type of speech feature (eg. "fbank"
        #: for log filterbank energies).
        self.feat_type = feat_type

        #: An arbitrary string representing the transcription tokenization
        #: used (eg. "phonemes", "tones", "joint", or "characters").
        self.label_type = label_type

        # Setting up directories
        # Set the directory names
        self.feat_dir = self.get_feat_dir()
        self.wav_dir = self.get_wav_dir()
        self.label_dir = self.get_label_dir()

        logger.debug("Setting up directories for this Corpus object at %s", self.tgt_dir)
        self.set_and_check_directories(self.tgt_dir)

        # Label-related stuff
        if labels is not None:
            self.labels = labels
            found_labels = determine_labels(self.tgt_dir, label_type)
            if found_labels != self.labels:
                raise LabelMismatchException("User specified labels, {}, do"
                    " not match those automatically found, {}.".format(labels,
                    found_labels))
        else:
            self.labels = determine_labels(self.tgt_dir, label_type)
        self.vocab_size = len(self.labels)
        self.LABEL_TO_INDEX, self.INDEX_TO_LABEL = self.initialize_labels(self.labels)
        logger.info("Corpus label set: \n\t{}".format(self.labels))

        # This is a lazy function that assumes wavs are already in the WAV dir
        # but only creates features if necessary
        logger.debug("Preparing features")
        self.prepare_feats()
        self._num_feats = None

        self.train_prefixes = [] # type: List[str]
        self.valid_prefixes = [] # type: List[str]
        self.test_prefixes = [] # type: List[str]
        # This is also lazy if the {train,valid,test}_prefixes.txt files exist.
        self.make_data_splits(max_samples=max_samples)

        # Sort the training prefixes by size for more efficient training
        logger.debug("Training prefixes")
        self.train_prefixes = utils.sort_by_size(
            self.feat_dir, self.train_prefixes, feat_type)

        # Ensure no overlap between training and test sets
        try:
            ensure_no_set_overlap(
                self.get_train_fns()[0],
                self.get_valid_fns()[0],
                self.get_test_fns()[0]
            )
        except PersephoneException:
            logger.error("Got overlap between train, valid and test data sets")
            raise

        untranscribed_from_file = self.get_untranscribed_prefixes()
        untranscribed_found = find_untranscribed_wavs(self.get_wav_dir(), self.get_label_dir(), self.label_type)

        self.untranscribed_prefixes = list(set(untranscribed_from_file) | set(untranscribed_found))

        # TODO Need to contemplate whether Corpus objects have Utterance
        # objects or # not. Some of the TestBKW tests currently rely on this
        # for testing.
        self.utterances = []  # type: List[Utterance]

        self.pickle()

    @classmethod
    def from_elan(cls: Type[CorpusT], org_dir: Path, tgt_dir: Path,
                  feat_type: str = "fbank", label_type: str = "phonemes",
                  *,
                  utterance_filter: Callable[[Utterance], bool] = None,
                  label_segmenter: Optional[LabelSegmenter] = None,
                  speakers: List[str] = None, lazy: bool = True,
                  tier_prefixes: Tuple[str, ...] = ("xv", "rf")) -> CorpusT:
        """ Construct a `Corpus` from ELAN files.

        Args:
            org_dir: A path to the directory containing the unpreprocessed
                data.
            tgt_dir: A path to the directory where the preprocessed data will
                be stored.
            feat_type: A string describing the input speech features. For
                       example, "fbank" for log Mel filterbank features.
            label_type: A string describing the transcription labels. For example,
                         "phonemes" or "tones".
            utterance_filter: A function that returns False if an utterance
                should not be included in the corpus and True otherwise. This
                can be used to remove undesirable utterances for training, such as
                codeswitched utterances.
            label_segmenter: An object that has an attribute `segment_labels`,
                which is creates new `Utterance` instances from old ones,
                by segmenting the tokens in their `text` attribute. Note,
                `LabelSegmenter` might be better as a function, the only issue
                is it needs to carry with it a list of labels. This could
                potentially be a function attribute.
            speakers: A list of speakers to filter for. If `None`, utterances
                from all speakers are included.
            tier_prefixes: A collection of strings that prefix ELAN tiers to
                filter for. For example, if this is `("xv", "rf")`, then tiers
                named "xv", "xv@Mark", "rf@Rose" would be extracted if they
                existed.

        """
        # This currently bails out if label_segmenter is not provided
        if not label_segmenter:
            raise ValueError("A label segmenter must be provided via label_segmenter")

        # In case path is supplied as a string, make it a Path
        if isinstance(tgt_dir, str):
            tgt_dir = Path(tgt_dir)

        # Read utterances from org_dir.
        utterances = elan.utterances_from_dir(org_dir,
                                              tier_prefixes=tier_prefixes)

        # Filter utterances based on some criteria (such as codeswitching).
        if utterance_filter:
            utterances = [utter for utter in utterances if utterance_filter(utter)]
        utterances = utterance.remove_duplicates(utterances)

        # Segment the labels in the utterances appropriately
        if label_segmenter:
            utterances = [label_segmenter.segment_labels(utter) for utter in utterances]

        # Remove utterances without transcriptions.
        utterances = utterance.remove_empty_text(utterances)

        # Remove utterances with exceptionally short wav_files that are too
        # short for CTC to work.
        utterances = utterance.remove_too_short(utterances)

        tgt_dir.mkdir(parents=True, exist_ok=True)

        # TODO A lot of these methods aren't ELAN-specific. preprocess.elan was
        # only used to get the utterances. There could be another Corpus
        # factory method that takes Utterance objects. the fromElan and
        # fromPangloss constructors could call this.

        # Writes the transcriptions to the tgt_dir/label/ dir
        utterance.write_transcriptions(utterances, (tgt_dir / "label"),
                               label_type, lazy=lazy)
        # Extracts utterance level WAV information from the input file.
        wav.extract_wavs(utterances, (tgt_dir / "wav"), lazy=lazy)

        corpus = cls(feat_type, label_type, tgt_dir,
                     labels=label_segmenter.labels, speakers=speakers)
        corpus.utterances = utterances
        return corpus

    def get_wav_dir(self) -> Path:
        return self.tgt_dir / "wav"

    def get_feat_dir(self) -> Path:
        return self.tgt_dir / "feat"

    def get_label_dir(self) -> Path:
        return self.tgt_dir / "label"

    @property
    def train_prefix_fn(self) -> Path:
        return self.tgt_dir / "train_prefixes.txt"

    @property
    def valid_prefix_fn(self) -> Path:
        return self.tgt_dir / "valid_prefixes.txt"

    @property
    def test_prefix_fn(self) -> Path:
        return self.tgt_dir / "test_prefixes.txt"

    def set_and_check_directories(self, tgt_dir: Path) -> None:
        """
        Make sure that the required directories exist in the target directory.
        set variables accordingly.
        """
        logger.info("Setting up directories for corpus in %s", tgt_dir)

        # Check directories exist.
        if not tgt_dir.is_dir():
            raise FileNotFoundError(
                "The directory {} does not exist.".format(tgt_dir))
        if not self.wav_dir.is_dir():
            raise PersephoneException(
                "The supplied path requires a 'wav' subdirectory.")
        self.feat_dir.mkdir(parents=True, exist_ok=True)
        if not self.label_dir.is_dir():
            raise PersephoneException(
                "The supplied path requires a 'label' subdirectory.")

    def initialize_labels(self, labels: Set[str]) -> Tuple[dict, dict]:
        """Create mappings from label to index and index to label"""
        logger.debug("Creating mappings for labels")

        label_to_index = {label: index for index, label in enumerate(
                                 ["pad"] + sorted(list(labels)))}
        index_to_label = {index: phn for index, phn in enumerate(
                                 ["pad"] + sorted(list(labels)))}

        return label_to_index, index_to_label

    def prepare_feats(self) -> None:
        """ Prepares input features"""

        logger.debug("Preparing input features")
        self.feat_dir.mkdir(parents=True, exist_ok=True)

        should_extract_feats = False
        for path in self.wav_dir.iterdir():
            if not path.suffix == ".wav":
                logger.info("Non wav file found in wav directory: %s", path)
                continue
            prefix = os.path.basename(os.path.splitext(str(path))[0])
            mono16k_wav_path = self.feat_dir / "{}.wav".format(prefix)
            feat_path = self.feat_dir / "{}.{}.npy".format(prefix, self.feat_type)
            if not feat_path.is_file():
                # Then we should extract feats
                should_extract_feats = True
                if not mono16k_wav_path.is_file():
                    feat_extract.convert_wav(path, mono16k_wav_path)

        # TODO Should be extracting feats on a per-file basis. Right now we
        # check if any feats files don't exist and then do all the feature
        # extraction.
        if should_extract_feats:
            feat_extract.from_dir(self.feat_dir, self.feat_type)

    def make_data_splits(self, max_samples: int) -> None:
        """ Splits the utterances into training, validation and test sets."""

        train_f_exists = self.train_prefix_fn.is_file()
        valid_f_exists = self.valid_prefix_fn.is_file()
        test_f_exists = self.test_prefix_fn.is_file()

        if train_f_exists and valid_f_exists and test_f_exists:
            logger.debug("Split for training, validation and tests specified by files")
            self.train_prefixes = self.read_prefixes(self.train_prefix_fn)
            self.valid_prefixes = self.read_prefixes(self.valid_prefix_fn)
            self.test_prefixes = self.read_prefixes(self.test_prefix_fn)
            return

        # Otherwise we now need to load prefixes for other cases addressed
        # below
        prefixes = self.determine_prefixes()
        prefixes = utils.filter_by_size(
            self.feat_dir, prefixes, self.feat_type, max_samples)

        if not train_f_exists and not valid_f_exists and not test_f_exists:
            logger.debug("No files supplied to define the split for training, validation"
                         " and tests. Using default.")
            train_prefixes, valid_prefixes, test_prefixes = self.divide_prefixes(prefixes)
            self.train_prefixes = train_prefixes
            self.valid_prefixes = valid_prefixes
            self.test_prefixes = test_prefixes
            self.write_prefixes(train_prefixes, self.train_prefix_fn)
            self.write_prefixes(valid_prefixes, self.valid_prefix_fn)
            self.write_prefixes(test_prefixes, self.test_prefix_fn)
        elif not train_f_exists and valid_f_exists and test_f_exists:
            # Then we just make all other prefixes training prefixes.
            self.valid_prefixes = self.read_prefixes(self.valid_prefix_fn)
            self.test_prefixes = self.read_prefixes(self.test_prefix_fn)
            train_prefixes = list(
                set(prefixes) - set(self.valid_prefixes))
            self.train_prefixes = list(
                set(train_prefixes) - set(self.test_prefixes))
            self.write_prefixes(self.train_prefixes, self.train_prefix_fn)
        else:
            raise NotImplementedError(
                "The following case has not been implemented:" +
                "{} exists - {}\n".format(self.train_prefix_fn, train_f_exists) +
                "{} exists - {}\n".format(self.valid_prefix_fn, valid_f_exists) +
                "{} exists - {}\n".format(self.test_prefix_fn, test_f_exists))

    @staticmethod
    def read_prefixes(prefix_fn: Path) -> List[str]:
        if not prefix_fn.is_file():
            logger.critical("Expected a prefix file at path {}, but this path is"
                           " not a file".format(prefix_fn))
        assert prefix_fn.is_file(), "Path {} was not a file".format(prefix_fn)
        with prefix_fn.open() as prefix_f:
            prefixes = [line.strip() for line in prefix_f]
        if not prefixes:
            raise PersephoneException(
                "Empty prefix file {}. Either delete it\
                or put something in it".format(prefix_fn))
        return prefixes

    @staticmethod
    def write_prefixes(prefixes: List[str], prefix_fn: Path) -> None:
        if not prefixes:
            raise PersephoneException(
                "No prefixes. Will not write {}".format(prefix_fn))
        with prefix_fn.open("w") as prefix_f:
            for prefix in prefixes:
                print(prefix, file=prefix_f)

    @staticmethod
    def divide_prefixes(prefixes: List[str], *, seed:int=0) -> Tuple[List[str], List[str], List[str]]:
        """Divide data into training, validation and test subsets"""
        if len(prefixes) < 3:
            raise PersephoneException(
                "{} cannot be split into 3 groups as it only has {} items".format(prefixes, len(prefixes))
            )
        Ratios = namedtuple("Ratios", ["train", "valid", "test"])
        ratios=Ratios(.90, .05, .05)
        train_end = int(ratios.train*len(prefixes))
        valid_end = int(train_end + ratios.valid*len(prefixes))

        # We must make sure that at least one element exists in test
        if valid_end == len(prefixes):
            valid_end -= 1

        # If train_end and valid_end are the same we end up with no valid_prefixes
        # so we must ensure at least one prefix is placed in this category
        if train_end == valid_end:
            train_end -= 1

        random.seed(seed)
        random.shuffle(prefixes)

        train_prefixes = prefixes[:train_end]
        valid_prefixes = prefixes[train_end:valid_end]
        test_prefixes = prefixes[valid_end:]

        assert train_prefixes, "Got empty set for training data"
        assert valid_prefixes, "Got empty set for validation data"
        assert test_prefixes, "Got empty set for testing data"

        return train_prefixes, valid_prefixes, test_prefixes

    def indices_to_labels(self, indices: Sequence[int]) -> List[str]:
        """ Converts a sequence of indices into their corresponding labels."""

        return [(self.INDEX_TO_LABEL[index]) for index in indices]

    def labels_to_indices(self, labels: Sequence[str]) -> List[int]:
        """ Converts a sequence of labels into their corresponding indices."""

        return [self.LABEL_TO_INDEX[label] for label in labels]

    @property
    def num_feats(self):
        """ The number of features per time step in the corpus. """
        if not self._num_feats:
            filename = self.get_train_fns()[0][0]
            feats = np.load(filename)
            # pylint: disable=maybe-no-member
            if len(feats.shape) == 3:
                # Then there are multiple channels of multiple feats
                self._num_feats = feats.shape[1] * feats.shape[2]
            elif len(feats.shape) == 2:
                # Otherwise it is just of shape time x feats
                self._num_feats = feats.shape[1]
            else:
                raise ValueError(
                    "Feature matrix of shape %s unexpected" % str(feats.shape))
        return self._num_feats

    def prefixes_to_fns(self, prefixes: List[str]) -> Tuple[List[str], List[str]]:
        """ Fetches the file paths to the features files and labels files
        corresponding to the provided list of features"""
        # TODO Return pathlib.Paths
        feat_fns = [str(self.feat_dir / ("%s.%s.npy" % (prefix, self.feat_type)))
                    for prefix in prefixes]
        label_fns = [str(self.label_dir / ("%s.%s" % (prefix, self.label_type)))
                      for prefix in prefixes]
        return feat_fns, label_fns

    def get_train_fns(self) -> Tuple[List[str], List[str]]:
        """ Fetches the training set of the corpus.

        Outputs a Tuple of size 2, where the first element is a list of paths
        to input features files, one per utterance. The second element is a list
        of paths to the transcriptions.
        """
        return self.prefixes_to_fns(self.train_prefixes)

    def get_valid_fns(self) -> Tuple[List[str], List[str]]:
        """ Fetches the validation set of the corpus."""
        return self.prefixes_to_fns(self.valid_prefixes)

    def get_test_fns(self) -> Tuple[List[str], List[str]]:
        """ Fetches the test set of the corpus."""
        return self.prefixes_to_fns(self.test_prefixes)

    def get_untranscribed_prefixes(self) -> List[str]:
        """
        The file "untranscribed_prefixes.txt" will specify prefixes which
        do not have an associated transcription file if placed in the target directory.

        This will fetch those prefixes from that file and will return an empty
        list if that file does not exist.

        See find_untranscribed_wavs function for finding untranscribed prefixes in an
        experiment directory.
        """
        return get_untranscribed_prefixes_from_file(self.tgt_dir)

    def get_untranscribed_fns(self) -> List[str]:
        feat_fns = [os.path.join(str(self.feat_dir), "%s.%s.npy" % (prefix, self.feat_type))
                    for prefix in self.untranscribed_prefixes]
        return feat_fns

    def determine_prefixes(self) -> List[str]:
        label_prefixes = [str(path.relative_to(self.label_dir).with_suffix(""))
                          for path in
                          self.label_dir.glob("**/*.{}".format(self.label_type))]
        wav_prefixes = [str(path.relative_to(self.wav_dir).with_suffix(""))
                          for path in
                          self.wav_dir.glob("**/*.{}".format("wav"))]

        # Take the intersection; sort for determinism.
        prefixes = sorted(list(set(label_prefixes) & set(wav_prefixes)))

        if not prefixes:
            raise PersephoneException("""WARNING: Corpus object has no data. Are you sure
            it's in the correct directories? WAVs should be in {} and
            transcriptions in {} with the extension .{}""".format(
                self.wav_dir, self.label_dir, self.label_type))

        return prefixes

    def review(self) -> None:
        """ Used to play the WAV files and compare with the transcription. """

        for prefix in self.determine_prefixes():
            print("Utterance: {}".format(prefix))
            wav_fn = self.feat_dir / "{}.wav".format(prefix)
            label_fn = self.label_dir / "{}.{}".format(prefix,self.label_type)
            with label_fn.open() as f:
                transcript = f.read().strip()
            print("Transcription: {}".format(transcript))
            subprocess.run(["play", str(wav_fn)])

    def pickle(self) -> None:
        """ Pickles the Corpus object in a file in tgt_dir. """

        pickle_path = self.tgt_dir / "corpus.p"
        logger.debug("pickling %r object and saving it to path %s", self, pickle_path)
        with pickle_path.open("wb") as f:
            pickle.dump(self, f)

    @classmethod
    def from_pickle(cls: Type[CorpusT], tgt_dir: Path) -> CorpusT:
        pickle_path = tgt_dir / "corpus.p"
        logger.debug("Creating Corpus object from pickle file path %s", pickle_path)
        with pickle_path.open("rb") as f:
            return pickle.load(f)

def determine_labels(target_dir: Path, label_type: str) -> Set[str]:
    """ Returns a set of all phonemes found in the corpus. Assumes that WAV files and
    label files are split into utterances and segregated in a directory which contains a
    "wav" subdirectory and "label" subdirectory.

    Arguments:
        target_dir: A `Path` to the directory where the corpus data is found
        label_type: The type of label we are creating the label set from. For example
                    "phonemes" would only search for labels for that type.
    """
    logger.info("Finding phonemes of type %s in directory %s", label_type, target_dir)

    label_dir = target_dir / "label/"
    if not label_dir.is_dir():
        raise FileNotFoundError(
            "The directory {} does not exist.".format(target_dir))

    phonemes = set() # type: Set[str]
    for fn in os.listdir(str(label_dir)):
        if fn.endswith(str(label_type)):
            with (label_dir / fn).open("r", encoding=ENCODING) as f:
                try:
                    line_phonemes = set(f.readline().split())
                except UnicodeDecodeError:
                    logger.error("Unicode decode error on file %s", fn)
                    print("Unicode decode error on file {}".format(fn))
                    raise
                phonemes = phonemes.union(line_phonemes)
    return phonemes


# coding: utf-8

import unittest
import httpretty
import json
from backlog.base import BacklogAPI

API_ENDPOINT = "https://{space}.backlog.jp/api/v2/{uri}"

class TestGit(unittest.TestCase):
    def setUp(self):
        self.api = BacklogAPI("space", "api_key")
        self.space = "space"
        self.api_key = "api_key"

    #
    # For Git Repository
    #
    @httpretty.activate
    def test_list_git_repositories(self):
        _project_id_or_key = 151
        _uri = "projects/{projectIdOrKey}/git/repositories".format(projectIdOrKey=_project_id_or_key)

        expects = [
            {
                "id": 1,
                "projectId": _project_id_or_key,
                "name": "app",
                "description": "",
                "hookUrl": None,
                "httpUrl": "https://xx.backlog.jp/git/BLG/app.git",
                "sshUrl": "xx@xx.git.backlog.jp:/BLG/app.git",
                "displayOrder": 0,
                "pushedAt": None,
                "createdUser": {
                    "id": 1,
                    "userId": "admin",
                    "name": "admin",
                    "roleType": 1,
                    "lang": "ja",
                    "mailAddress": "eguchi@nulab.example"
                },
                "created": "2013-05-30T09:11:36Z",
                "updatedUser": {
                    "id": 1,
                    "userId": "admin",
                    "name": "admin",
                    "roleType": 1,
                    "lang": "ja",
                    "mailAddress": "eguchi@nulab.example"
                },
                "updated": "2013-05-30T09:11:36Z"
            }
        ]

        httpretty.register_uri(
            httpretty.GET,
            API_ENDPOINT.format(space=self.space, uri=_uri),
            body=json.dumps(expects),
            status=200
        )

        resp = self.api.git.list_git_repositories(projectIdOrKey=_project_id_or_key)

        self.assertEqual(expects, resp)

    @httpretty.activate
    def test_get_git_repositories(self):
        pass

    #
    # For Pull Request
    #
    @httpretty.activate
    def test_list_pull_requests(self):
        _project_id_or_key = 3
        _repo_id_or_name = 5
        _uri = "projects/{projectIdOrKey}/git/repositories/{repoIdOrName}/pullRequests".format(
            projectIdOrKey=_project_id_or_key,
            repoIdOrName=_repo_id_or_name
        )

        expects = [
            {
                "id": 2,
                "projectId": 3,
                "repositoryId": 5,
                "number": 1,
                "summary": "test",
                "description": "test data",
                "base": "master",
                "branch": "develop",
                "status": {
                    "id": 1,
                    "name": "Open"
                },
                "assignee": {
                    "id": 5,
                    "userId": "testuser2",
                    "name": "testuser2",
                    "roleType": 1,
                    "lang": None,
                    "mailAddress": "testuser2@nulab.test"
                },
                "issue": {
                    "id": 31
                },
                "baseCommit": None,
                "branchCommit": None,
                "closeAt": None,
                "mergeAt": None,
                "createdUser": {
                    "id": 1,
                    "userId": "admin",
                    "name": "admin",
                    "roleType": 1,
                    "lang": "ja",
                    "mailAddress": "eguchi@nulab.example"
                },
                "created": "2015-04-23T03:04:14Z",
                "updatedUser": {
                    "id": 1,
                    "userId": "admin",
                    "name": "admin",
                    "roleType": 1,
                    "lang": "ja",
                    "mailAddress": "eguchi@nulab.example"
                },
                "updated": "2015-04-23T03:04:14Z"
            }
        ]

        httpretty.register_uri(
            httpretty.GET,
            API_ENDPOINT.format(space=self.space, uri=_uri),
            body=json.dumps(expects),
            status=200
        )

        resp = self.api.git.list_pull_requests(
            projectIdOrKey=_project_id_or_key,
            repoIdOrName=_repo_id_or_name
        )

        self.assertEqual(expects, resp)

    @httpretty.activate
    def test_count_pull_requests(self):
        pass

    @httpretty.activate
    def test_add_pull_request(self):
        _project_id_or_key = 3
        _repository_id_or_key = 5
        _uri = "projects/{projectIdOrKey}/git/repositories/{repoIdOrName}/pullRequests".format(
            projectIdOrKey=_project_id_or_key,
            repoIdOrName=_repository_id_or_key
        )

        expects = {
            "id": 2,
            "projectId": _project_id_or_key,
            "repositoryId": _repository_id_or_key,
            "number": 1,
            "summary": "test",
            "description": "test data",
            "base": "master",
            "branch": "develop",
            "status": {
                "id": 1,
                "name": "Open"
            },
            "assignee": {
                "id": 5,
                "userId": "testuser2",
                "name": "testuser2",
                "roleType": 1,
                "lang": None,
                "mailAddress": "testuser2@nulab.test"
            },
            "issue": {
                "id": 1,
                "projectId": 1,
                "issueKey": "BLG-1",
                "keyId": 1,
                "issueType": {
                    "id": 2,
                    "projectId": 1,
                    "name": "タスク",
                    "color": "#7ea800",
                    "displayOrder": 0
                },
                "summary": "first issue",
                "description": "",
                "resolutions": None,
                "priority": {
                    "id": 3,
                    "name": "中"
                },
                "status": {
                    "id": 1,
                    "name": "未対応"
                },
                "assignee": {
                    "id": 2,
                    "name": "eguchi",
                    "roleType": 2,
                    "lang": None,
                    "mailAddress": "eguchi@nulab.example"
                },
                "category": [],
                "versions": [],
                "milestone": [
                    {
                        "id": 30,
                        "projectId": 1,
                        "name": "wait for release",
                        "description": "",
                        "startDate": None,
                        "releaseDueDate": None,
                        "archived": False
                    }
                ],
                "startDate": None,
                "dueDate": None,
                "estimatedHours": None,
                "actualHours": None,
                "parentIssueId": None,
                "createdUser": {
                    "id": 1,
                    "userId": "admin",
                    "name": "admin",
                    "roleType": 1,
                    "lang": "ja",
                    "mailAddress": "eguchi@nulab.example"
                },
                "created": "2012-07-23T06:10:15Z",
                "updatedUser": {
                    "id": 1,
                    "userId": "admin",
                    "name": "admin",
                    "roleType": 1,
                    "lang": "ja",
                    "mailAddress": "eguchi@nulab.example"
                },
                "updated": "2013-02-07T08:09:49Z",
                "customFields": [],
                "attachments": [
                    {
                        "id": 1,
                        "name": "IMGP0088.JPG",
                        "size": 85079
                    }
                ],
                "sharedFiles": [
                    {
                        "id": 454403,
                        "type": "file",
                        "dir": "/ユーザアイコン/",
                        "name": "01_サラリーマン.png",
                        "size": 2735,
                        "createdUser": {
                            "id": 5686,
                            "userId": "takada",
                            "name": "takada",
                            "roleType": 2,
                            "lang": "ja",
                            "mailAddress": "takada@nulab.example"
                        },
                        "created": "2009-02-27T03:26:15Z",
                        "updatedUser": {
                            "id": 5686,
                            "userId": "takada",
                            "name": "takada",
                            "roleType": 2,
                            "lang": "ja",
                            "mailAddress": "takada@nulab.example"
                        },
                        "updated": "2009-03-03T16:57:47Z"
                    }
                ],
                "stars": [
                    {
                        "id": 10,
                        "comment": None,
                        "url": "https://xx.backlog.jp/view/BLG-1",
                        "title": "[BLG-1] first issue | 課題の表示 - Backlog",
                        "presenter": {
                            "id": 2,
                            "userId": "eguchi",
                            "name": "eguchi",
                            "roleType": 2,
                            "lang": "ja",
                            "mailAddress": "eguchi@nulab.example"
                        },
                        "created": "2013-07-08T10:24:28Z"
                    }
                ]
            },
            "baseCommit": None,
            "branchCommit": None,
            "closeAt": None,
            "mergeAt": None,
            "createdUser": {
                "id": 1,
                "userId": "admin",
                "name": "admin",
                "roleType": 1,
                "lang": "ja",
                "mailAddress": "eguchi@nulab.example"
            },
            "created": "2015-04-23T03:04:14Z",
            "updatedUser": {
                "id": 1,
                "userId": "admin",
                "name": "admin",
                "roleType": 1,
                "lang": "ja",
                "mailAddress": "eguchi@nulab.example"
            },
            "updated": "2015-04-23T03:04:14Z"
        }

        httpretty.register_uri(
            httpretty.POST,
            API_ENDPOINT.format(space=self.space, uri=_uri),
            body=json.dumps(expects),
            status=200
        )

        resp = self.api.git.add_pull_request(
            projectIdOrKey=_project_id_or_key,
            repoIdOrName=_repository_id_or_key)

        self.assertEqual(expects, resp)

    @httpretty.activate
    def test_get_pull_requests(self):
        pass

    @httpretty.activate
    def test_update_pull_requests(self):
        pass

    @httpretty.activate
    def test_get_pull_request_comment(self):
        pass

    @httpretty.activate
    def test_add_pull_request_comment(self):
        _project_id_or_key = 3
        _repository_id_or_key = 5
        _number = 1
        _content = "test",
        _notified_user_ids = []
        _uri = "projects/{projectIdOrKey}/git/repositories/{repoIdOrName}/pullRequests/{number}/comments".format(
            projectIdOrKey=_project_id_or_key,
            repoIdOrName=_repository_id_or_key,
            number=_number
        )

        expects = {
            "id": 35,
            "content": "from api",
            "changeLog": [
                {
                    "field": "dependentIssue",
                    "newValue": "GIT-3",
                    "originalValue": None
                }
            ],
            "createdUser": {
                "id": 1,
                "userId": "admin",
                "name": "admin",
                "roleType": 1,
                "lang": "ja",
                "mailAddress": "eguchi@nulab.example"
            },
            "created": "2015-05-14T01:53:38Z",
            "updated": "2015-05-14T01:53:38Z",
            "stars": [],
            "notifications": []
        }

        httpretty.register_uri(
            httpretty.POST,
            API_ENDPOINT.format(space=self.space, uri=_uri),
            body=json.dumps(expects),
            status=200
        )

        resp = self.api.git.add_pull_request_comment(
            projectIdOrKey=_project_id_or_key,
            repoIdOrName=_repository_id_or_key,
            number=_number,
            content=_content,
            notifiedUserIds=_notified_user_ids
        )

        self.assertEqual(expects, resp)

    @httpretty.activate
    def test_count_pull_request_comments(self):
        pass

    @httpretty.activate
    def test_update_pull_request_comment(self):
        pass

    @httpretty.activate
    def test_list_pull_request_attachments(self):
        pass

    @httpretty.activate
    def test_get_pull_request_attachment(self):
        pass

    @httpretty.activate
    def test_delelte_pull_request_attachment(self):
        pass

if __name__ == "__main__":
    unittest.main(warnings='ignore')


from dataclasses import dataclass
from unittest import mock

import pytest

from genie.backends.database.backend import DatabaseBackend
from genie.backends.database.models import BaseModel

@dataclass
class TestModel(BaseModel):
    table_name = 'Test'
    name: str

class TestBaseModel:

    @pytest.fixture
    def model(self):
        return TestModel(name='test')

    def test_connection_returns_instance(self, model):
        assert isinstance(model.connection(), DatabaseBackend)

    def test_to_dict_calls_asdict(self, model):
        with mock.patch(
            'genie.backends.database.models.asdict'
        ) as mock_asdict:
            model.to_dict()

        mock_asdict.assert_called_with(model)

    def test_get_calls_connection(self, model):
        with mock.patch(
            'genie.backends.database.backend.DatabaseBackend.get'
        ) as connection:
            TestModel.get(a='a', b='b')

        connection.assert_called_with(a='a', b='b')

    def test_delete_calls_connection(self, model):
        with mock.patch(
            'genie.backends.database.backend.DatabaseBackend.delete'
        ) as connection:
            TestModel.delete(a='a')

        connection.assert_called_with(a='a')

    def test_save_calls_connection(self, model):
        with mock.patch(
            'genie.backends.database.backend.DatabaseBackend.save'
        ) as connection:
            model.save()

        connection.assert_called_with(**model.to_dict())


from sklearn import linear_model
import numpy as np
import matplotlib.pyplot as plt
np.random.seed(1)

# 第一步 生成数据
X = np.random.rand(20, 1) * 10 # sklearn线性模型只接受二维数据
y = 3 * X + (5 + np.random.randn(20, 1))

# 第二步 定义模型
model = linear_model.LinearRegression()

# 第三步 拟合数据
model.fit(X, y)

# 第四步 预测（以下两种方法都可以）
# y_pred = k * X + b
y_pred = model.predict(X)

# 第四步 展示拟合结果
b = model.intercept_.item()
k = model.coef_.item()
print('截距为：', b)  #截距
print('斜率为：', k)  #线性模型的系数

plt.figure(figsize = (12,12))
plt.scatter(X.reshape(-1), y.reshape(-1))
plt.plot(X.reshape(-1), y_pred.reshape(-1), 'r-', lw=5)
plt.tick_params(labelsize=20)
plt.xlabel('X', fontsize = 25)
plt.ylabel('y', fontsize = 25)
plt.xlim(0, 10)
plt.ylim(8, 35)
plt.title('Sklearn result of linear regression', fontsize = 30)
plt.show()


from libgravatar import Gravatar
from sqlalchemy.orm import Session

from src.database.models import User
from src.schemas import UserModel

async def get_user_by_email(email: str, db: Session) -> User | None:
    """
    Get user by e-mail.

    :param email: User's e-mail.
    :type email: str
    :param db: Database session.
    :type db: Session
    :return: Return the found user, or None if it does not exist
    :rtype: User | None
    """
    return db.query(User).filter(User.email == email).first()

async def create_user(body: UserModel, db: Session) -> User:
    """
    Create user.

    :param body: User body.
    :type body: UserModel
    :param db: Database session.
    :type db: Session
    :return: Created user.
    :rtype: User
    """
    avatar = None
    try:
        g = Gravatar(body.email)
        avatar = g.get_image()
    except Exception as e:
        print(e)
    new_user = User(**body.dict(), avatar=avatar)
    db.add(new_user)
    db.commit()
    db.refresh(new_user)
    return new_user

async def update_token(user: User, token: str | None, db: Session) -> None:
    """
    Update user token

    :param user: Current user.
    :type user: User
    :param token: User's token.
    :type token: str
    :param db: Database session.
    :type db: Session
    :return: Not return result.
    """
    user.refresh_token = token
    db.commit()

async def confirmed_email(email: str, db: Session) -> None:
    """
    Set user's e-mail confirmed.

    :param email: User's e-mail
    :type email: str
    :param db: Database session.
    :type db: Session
    :return: Not return result.
    """
    user = await get_user_by_email(email, db)
    user.confirmed = True
    db.commit()

async def update_avatar(email, url: str, db: Session) -> User:
    """
    Update user's avatar

    :param email: User's e-mail.
    :type email: str
    :param url: User avatar link.
    :type url: str
    :param db: Database session.
    :type db: Session
    :return: Updated user.
    :rtype: User

    """
    user = await get_user_by_email(email, db)
    user.avatar = url
    db.commit()
    return user


import cv2 as cv
import numpy as np

print(np.__version__)

train_data = np.load('training_data.npy')

for data in train_data:
    img = data[0]
    key_press = data[1]
    cv.imshow('test', img)
    print(key_press)
    if cv.waitKey(25) & 0xFF == ord('q'):
        cv.destroyAllWindows()
        break


from collections import defaultdict

class InvertedIndex:
    def __init__(self):
        self.index = defaultdict(list)
        self.document_metadata = {}

    def add_document(self, document_id, content, metadata=None):
        tokens = content.lower().split()  # Convert to lowercase for case-insensitive search
        for token in tokens:
            self.index[token].append(document_id)
        if metadata:
            self.document_metadata[document_id] = metadata

    def search_documents(self, keywords):
        keyword_set = set(keywords)
        matching_documents = defaultdict(int)
        for keyword in keyword_set:
            if keyword in self.index:
                for document_id in self.index[keyword]:
                    matching_documents[document_id] += 1
        sorted_results = sorted(matching_documents.items(), key=lambda x: x[1], reverse=True)
        return [result[0] for result in sorted_results]

    def remove_document(self, document_id):
        for postings in self.index.values():
            if document_id in postings:
                postings.remove(document_id)
                if not postings:
                    del postings
        if document_id in self.document_metadata:
            del self.document_metadata[document_id]

index = InvertedIndex()

index.add_document(1, "Apple Inc. announces new product lineup.", {"date": "2023-06-15", "source": "PR Newswire"})
index.add_document(2, "Google reports strong financial results for Q2.",
                   {"date": "2023-06-16", "source": "Business Insider"})
index.add_document(3, "Amazon launches new subscription service.", {"date": "2023-06-14", "source": "Reuters"})
index.add_document(4, "Microsoft introduces new cloud computing solution.",
                   {"date": "2023-06-17", "source": "TechCrunch"})
index.add_document(5, "Tesla unveils new electric vehicle model.", {"date": "2023-06-17", "source": "CNN"})
index.add_document(6, "Facebook announces expansion into virtual reality market.",
                   {"date": "2023-06-16", "source": "The Verge"})
index.add_document(7, "Samsung releases latest smartphone with advanced features.",
                   {"date": "2023-06-15", "source": "Gizmodo"})
index.add_document(8, "Netflix reports record-breaking subscriber growth.", {"date": "2023-06-18", "source": "Variety"})
index.add_document(9, "IBM introduces new AI-powered analytics platform.", {"date": "2023-06-18", "source": "Forbes"})
index.add_document(10, "General Electric announces major renewable energy project.",
                   {"date": "2023-06-19", "source": "CNBC"})
index.add_document(11, "Toyota unveils all-electric SUV with impressive range.",
                   {"date": "2023-06-19", "source": "The Guardian"})
index.add_document(12, "Intel launches revolutionary processor for high-performance computing.",
                   {"date": "2023-06-20", "source": "PCMag"})
index.add_document(13, "SpaceX successfully launches satellite into orbit.",
                   {"date": "2023-06-20", "source": "Space.com"})
index.add_document(14, "Coca-Cola introduces new line of healthy beverages.",
                   {"date": "2023-06-21", "source": "Bloomberg"})
index.add_document(15, "Adobe releases latest version of creative software suite.",
                   {"date": "2023-06-21", "source": "Engadget"})
index.add_document(16, "Walmart announces expansion into online grocery delivery.",
                   {"date": "2023-06-22", "source": "TechRadar"})
index.add_document(17, "Volkswagen unveils electric vehicle charging infrastructure plan.",
                   {"date": "2023-06-22", "source": "Reuters"})
index.add_document(18, "Twitter introduces new feature to combat misinformation.",
                   {"date": "2023-06-23", "source": "The New York Times"})
index.add_document(19, "JPMorgan Chase reports strong quarterly earnings.",
                   {"date": "2023-06-23", "source": "Financial Times"})
index.add_document(20, "Boeing showcases next-generation aircraft with advanced technologies.",
                   {"date": "2023-06-24", "source": "Aviation Week"})


import random

import cv2
import mujoco_py
import numpy as np
import warnings
from PIL import Image
from gym.spaces import Box, Dict

from multiworld.core.multitask_env import MultitaskEnv
from multiworld.core.wrapper_env import ProxyEnv
from multiworld.envs.env_util import concatenate_box_spaces
from multiworld.envs.env_util import get_stat_in_paths, create_stats_ordered_dict

from gym.spaces import Box, Dict

class GymToMultiEnv(ProxyEnv): # MultitaskEnv):
    def __init__(
            self,
            wrapped_env,
    ):
        """Minimal env to convert a gym env to one with dict observations"""
        self.quick_init(locals())
        super().__init__(wrapped_env)

        obs_box = wrapped_env.observation_space
        self.observation_space = Dict([
            ('observation', obs_box),
            ('state_observation', obs_box),
        ])

    def step(self, action):
        obs, reward, done, info = self.wrapped_env.step(action)
        new_obs = dict(
            observation=obs,
            state_observation=obs,
        )
        return new_obs, reward, done, info

    def reset(self):
        obs = self.wrapped_env.reset()
        new_obs = dict(
            observation=obs,
            state_observation=obs,
        )
        return new_obs

    def _get_obs(self):
        raise NotImplementedError()

class MujocoGymToMultiEnv(GymToMultiEnv):
    def set_state(self, qpos, qvel):
        assert qpos.shape == (self.model.nq,) and qvel.shape == (self.model.nv,)
        old_state = self.sim.get_state()
        new_state = mujoco_py.MjSimState(old_state.time, qpos, qvel,
                                         old_state.act, old_state.udd_state)
        self.sim.set_state(new_state)
        self.sim.forward()

    @property
    def dt(self):
        return self.model.opt.timestep * self.frame_skip

    def do_simulation(self, ctrl, n_frames=None):
        if n_frames is None:
            n_frames = self.frame_skip
        if self.sim.data.ctrl is not None and ctrl is not None:
            self.sim.data.ctrl[:] = ctrl
        for _ in range(n_frames):
            self.sim.step()

    def render(self, mode='human'):
        if mode == 'rgb_array':
            self._get_viewer().render()
            # window size used for old mujoco-py:
            width, height = 500, 500
            width, height = 4000, 4000
            data = self._get_viewer().read_pixels(width, height, depth=False)
            # original image is upside-down, so flip it
            return data[::-1, :, :]
        elif mode == 'human':
            self._get_viewer().render()

    def close(self):
        if self.viewer is not None:
            self.viewer.finish()
            self.viewer = None

    def _get_viewer(self):
        if self.viewer is None:
            self.viewer = mujoco_py.MjViewer(self.sim)
            self.viewer_setup()
        return self.viewer

    def get_body_com(self, body_name):
        return self.data.get_body_xpos(body_name)

    def state_vector(self):
        return np.concatenate([
            self.sim.data.qpos.flat,
            self.sim.data.qvel.flat
        ])

    def get_image(self, width=84, height=84, camera_name=None):
        return self.sim.render(
            width=width,
            height=height,
            camera_name=camera_name,
        )

    def initialize_camera(self, init_fctn):
        sim = self.sim
        viewer = mujoco_py.MjRenderContextOffscreen(sim, device_id=self.device_id)
        init_fctn(viewer.cam)
        sim.add_render_context(viewer)

    def get_diagnostics(self, paths, **kwargs):
        return {}