"""
Spokestack-Lite Speech Synthesizer

This module contains the SpeechSynthesizer class used to convert text to speech
using local TTS models trained on the Spokestack platform. A SpeechSynthesizer
instance can be passed to the TextToSpeechManager for playback.

Example:
    This example assumes that a TTS model was downloaded from the Spokestack
    platform and extracted to the :code:`model` directory. ::

        from spokestack.io.pyaudio import PyAudioOutput
        from spokestack.tts.manager import TextToSpeechManager, FORMAT_PCM16
        from spokestack.tts.lite import SpeechSynthesizer, BLOCK_LENGTH, SAMPLE_RATE

        tts = TextToSpeechManager(
            SpeechSynthesizer("./model"),
            PyAudioOutput(sample_rate=SAMPLE_RATE, frames_per_buffer=BLOCK_LENGTH),
            format_=FORMAT_PCM16)

        tts.synthesize("Hello world!")

"""

import importlib
import json
import os
import re
import typing as T
from collections import defaultdict

import numpy as np

from spokestack.models.tensorflow import TFLiteModel

# signal configuration
SAMPLE_RATE = 24000
HOP_LENGTH = 240
ENCODER_PAD = -2
BREAK_LENGTH = 0.1

# streaming/cross-fading configuration
FRAME_LENGTH = 63
FRAME_OVERLAP = 1
BLOCK_LENGTH = FRAME_LENGTH * HOP_LENGTH
BLOCK_OVERLAP = FRAME_OVERLAP * HOP_LENGTH
FADE_OUT = np.linspace(1, 0, BLOCK_OVERLAP, dtype=np.float32)
FADE_IN = FADE_OUT[::-1]

class SpeechSynthesizer:
    """
    Initialize a new lightweight speech synthesizer

    Args:
        model_path (str): Path to the extracted TTS model downloaded from the
            Spokestack platform

    """

    def __init__(self, model_path: str):
        # load NLP configuration
        self._lexicon = _load_lexicon(os.path.join(model_path, "lexicon.txt"))

        with open(os.path.join(model_path, "metadata.json")) as file:
            metadata = json.load(file)

        lang = metadata["language"]
        self._sym_to_id = {s: i for i, s in enumerate(metadata["alphabet"])}
        self._language: T.Any = importlib.import_module(f"spokestack.tts.lite.{lang}")
        self._nlp = self._language.nlp()

        # load the TTS models
        self._aligner = TFLiteModel(os.path.join(model_path, "align.tflite"))
        self._encoder = TFLiteModel(os.path.join(model_path, "encode.tflite"))
        self._decoder = TFLiteModel(os.path.join(model_path, "decode.tflite"))
        self._aligner_input_index = self._aligner.input_details[0]["index"]
        self._encoder_input_index = self._encoder.input_details[0]["index"]

    def synthesize(
        self, utterance: str, *_args: T.List, **_kwargs: T.Dict
    ) -> T.Iterator[np.array]:
        """
        Synthesize a text utterance to speech audio

        Args:
            utterance (str): The text string to synthesize

        Returns:
            Iterator[np.array]: A generator for returns a sequence of
            PCM-16 numpy audio blocks for playback, storage, etc.

        """

        # segment sentences into a list of phoneme/grapheme lists
        for tokens in self._parse(utterance):
            # convert tokens to a vector of ids
            inputs = self._vectorize(tokens)

            # run the aligner model
            self._aligner.resize(self._aligner_input_index, inputs.shape)
            inputs = self._aligner(inputs)[0]

            # run the encoder model
            self._encoder.resize(self._encoder_input_index, inputs.shape)
            encoded = self._encoder(inputs)[0]

            # stream the decoder model and cross-fade the output audio
            overlap = np.zeros([BLOCK_OVERLAP], dtype=np.float32)
            for i in range(FRAME_OVERLAP, len(encoded), FRAME_LENGTH):
                # decode the current frame, padding as need to fill the decoder's input
                inputs = encoded[i - FRAME_OVERLAP : i + FRAME_LENGTH]
                inputs = np.pad(
                    inputs,
                    [(0, (FRAME_LENGTH + FRAME_OVERLAP) - len(inputs)), (0, 0)],
                    "constant",
                    constant_values=ENCODER_PAD,
                )
                outputs = self._decoder(inputs)[0]

                # fade in the new block, convert to int16 and return it
                overlap += outputs[:BLOCK_OVERLAP] * FADE_IN
                block = np.hstack([overlap, outputs[BLOCK_OVERLAP:-BLOCK_OVERLAP]])
                yield (block * (2 ** 15 - 1)).astype(np.int16)

                # fade out the previous block for mixing with the next block
                overlap = outputs[-BLOCK_OVERLAP:] * FADE_OUT

            # add a break after each segment
            yield np.zeros([int(BREAK_LENGTH * SAMPLE_RATE)], dtype=np.int16)

    def _parse(self, text: str) -> T.Iterator[str]:
        # perform language-specific number conversions, abbreviation expansions, etc.
        text = self._language.clean(text)

        # escape characters used for phonetic substitution
        text = re.sub(r"{", "[", text)
        text = re.sub(r"}", "]", text)

        # segment and tokenize the text, and convert words to their phonetic
        # representations using the attached lexicon
        for sentence in self._nlp(text).sents:
            tokens = []
            for token in sentence:
                if token.pos_ in ["SYM", "PUNCT"]:
                    tokens.append(token.text_with_ws)
                else:
                    entry = self._lexicon.get(token.text.lower(), {})
                    ipa = entry.get(token.tag_, entry.get(None))
                    tokens.append(
                        f"{{{ipa}}}{token.whitespace_}" if ipa else token.text_with_ws
                    )
            yield re.sub(r"}\s+{", " ", "".join(tokens))

    def _vectorize(self, text: str) -> np.array:
        # start with bos token
        vector = [self._sym_to_id["^"]]

        while text:
            # check for curly braces and treat their contents as ipa
            matches = re.match(r"(.*?)\{(.+?)\}(.*)", text)

            # no ipa in this block, vectorize graphemes
            if not matches:
                vector.extend(self._vectorize_text(text))
                break

            # ipa found, vectorize leading text, then phones
            vector.extend(self._vectorize_text(matches.group(1)))
            vector.extend(self._vectorize_phones(matches.group(2)))
            text = matches.group(3)

        # append eos token
        vector.append(self._sym_to_id["~"])
        return np.array(vector, dtype=np.int32)

    def _vectorize_text(self, text: T.Union[str, T.List[str]]) -> T.List[int]:
        return [
            self._sym_to_id[c] for c in text if c in self._sym_to_id and c not in "_^~"
        ]

    def _vectorize_phones(self, phones: str) -> T.List[int]:
        return self._vectorize_text([f"@{c}" if c != " " else c for c in phones])

def _load_lexicon(path: str) -> T.Dict[str, T.Dict[T.Optional[str], str]]:
    lexicon: T.Dict[str, T.Dict[T.Optional[str], str]] = defaultdict(dict)

    with open(path, "r") as file:
        for line in file:
            # parse the the lexicon entry, discard any alternative pronunciations
            parts = line.strip().split("\t")
            if len(parts) > 1:
                word = parts[0].lower()
                ipa = parts[1].split(",")[0].strip()
                pos = parts[2] if len(parts) > 2 else None
                lexicon[word][pos] = ipa

    return lexicon


"""
    Run all PyDSTool tests and examples
"""
import os, time
from numpy import any

# Set this to be the command you want to use to invoke python
# in the tests. If None, then the $PYTHON environment variable will
# be checked first, and if empty 'python' will be used
pythonprogram = None

##### Select which sets of tests to run
# Test lists have entries on separate lines to make it easier to
# comment out the ones you don't want
#
# Basic PyDSTool modules
test_general = True
# Map generators
test_maps = True
# ODE integration with vode, no external compiler needed
test_vode = True
# Parameter estimation module
test_param_est = True
# Parameter estimation, requires external C compiler
test_param_est_C = True
# Symbolic differentiation module
test_symbolic = True
# Dopri integration; requires external C compiler
test_dopri = True
# Radau stiff integration; requires external C and fortran compilers
test_radau = True
# Continuation; no external compilers needed
test_pycont = True
# Continuation with auto; requires external compilers
test_pycont_auto = True

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

# Basic PyDSTool modules
general_list = [
    'test_variable_traj',
    'traj_gt0_test',
    'ModelSpec_test',
    'interp_piecewise_test',
    'objectdelete_test',
    ]

# Map generators
map_list = [
    'SLIP_2D_maps'
    ]

# ODE integration with vode, no external compiler needed
vode_list = [
    'poly_interp_test',
    'vode_withJac_test',
    'interp_vode_test',
    'fingermodel_vode',
    'fingermodel_auxvartest',
    'test_hybrid_extinputs',
    'HH_model',
    'HH_model_testbounds',
    'HH_loaded',
    'IF_model_test',
    'IF_squarespike_model',
    'IF_delaynet',
    'forced_spring'
    ]

# Parameter estimation module
param_est_list = [
    'joe_pest',
    'pest_test1',
    'pest_test2',
    'pest_test3',
    ]

# Parameter estimation, requires external C compiler
param_est_C_list = [
    'pest_test3_Cintegrator',
    'pest_test4_Cintegrator'
    ]

# Symbolic differentiation module
symbolic_list = [
    'Tutorial_SymbolicJac'
    ]

# Dopri integration; requires external C compiler
dopri_list = [
    'imprecise_event_test',
    'interp_dopri_test',
    'HH_model_Cintegrator',
    'HH_loaded_Cintegrator',
    'IF_delaynet_syn',
    'CIN',
    'HH_model_Cintegrator_testbounds',
    'Dopri_backwards_test'
    ]

# Radau stiff integration; requires external C and fortran compilers
radau_list = [
    'test_hybrid_extinputs_Cintegrator',
    'SLIP_2D_pdc',
    'DAE_example',
    'freefinger_noforce_radau',
    'sloppycell_example'
    ]

# Continuation; no external compilers needed
pycont_list = [
    'PyCont_Brusselator',
    'PyCont_Catalytic',
    'PyCont_ABCReaction',
    'PyCont_DiscPredPrey',
    'PyCont_Hopfield',
    'PyCont_hybrid_osc',
    'PyCont_LevelCurve',
    'PyCont_Logistic',
    'PyCont_NewLorenz',
    'PyCont_PredPrey',
    'PyCont_SaddleNode',
    ]

# Continuation with auto; requires external compilers
pycont_auto_list = [
    'PyCont_MorrisLecar_TypeI',
    'PyCont_MorrisLecar_TypeII',
    'PyCont_LPNeuron',
    'PyCont_HindmarshRose',
    'PyCont_ABReaction',
    'PyCont_Hamiltonian',
    'PyCont_Lorenz',
    'PyCont_vanDerPol',
    ]

do_external = any([test_dopri, test_radau, test_param_est_C, test_pycont_auto])

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

res = []
failed = []

def test(flist, whichpy=None, infostr=""):
    if whichpy is None or not whichpy:
        whichpy = os.getenv(key='PYTHON', default='python')

    failure = False
    for f in flist:
        fname = f+'.py'
        print("\n***** Testing script %s ****************************\n"%fname)
        try:
            e=os.system(whichpy + ' ' + fname)
        except:
            print("\n      Testing failed on test file %s"%fname)
            failed.append(fname)
            if not failure:
                res.append("%s: appears to be broken on your system"%infostr)
                failure = True
        else:
            if e in [0,3]:
                print("\n      Testing passed on test file %s"%fname)
            else:
                print("\n      Testing failed on test file %s"%fname)
                failed.append(fname)
                failure = True
        time.sleep(2)
    if failure:
        res.append("%s: appears to be broken on your system"%infostr)
    else:
        res.append("%s: appears to work on your system"%infostr)

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

print("***** Running all tests in order...\n")
print("Note: Depending on your settings, you may have to close matplotlib windows by hand in order to continue to the next test script\n")

if test_general:
    print("Testing general PyDSTool functions...\n")
    test(general_list, pythonprogram, "Basic PyDSTool functions")
else:
    res.append("Tests of basic PyDSTool functions: SKIPPED")

if test_maps:
    print("Testing map modules...\n")
    test(map_list, pythonprogram, "Map related modules")
else:
    res.append("Tests of map related modules: SKIPPED")

if test_vode:
    print("Testing modules using VODE integrator...\n")
    test(vode_list, pythonprogram, "VODE related modules")
else:
    res.append("Tests of VODE related modules: SKIPPED")

if test_symbolic:
    print("Testing symbolic differentiation module...\n")
    test(symbolic_list, pythonprogram, "Symbolic differentiation module")
else:
    res.append("Tests of symbolic differentiation module: SKIPPED")

if test_param_est:
    print("Testing parameter estimation module, no C compiler dependence...\n")
    test(param_est_list, pythonprogram, "Parameter estimation module")
else:
    res.append("Tests of parameter estimation module: SKIPPED")

if test_pycont:
    print("Testing PyCont module, no external compiler dependence...\n")
    test(pycont_list, pythonprogram, "PyCont")
else:
    res.append("Tests of PyCont with no external compiler dependence: SKIPPED")

if do_external:
    print("\n***** Now running tests that use external compilers...\n")

if test_dopri:
    print("Testing dopri integration; external C compiler dependence...\n")
    test(dopri_list, pythonprogram, "Dopri ODE systems")
else:
    res.append("Tests of Dopri ODE systems: SKIPPED")

if test_radau:
    print("Testing radau integration; external C, fortran compiler dependence...\n")
    test(radau_list, pythonprogram, "Radau ODE systems")
else:
    res.append("Tests of Radau ODE systems: SKIPPED")

if test_param_est_C:
    print("Testing parameter estimation module; with C compiler dependence...\n")
    test(param_est_C_list, pythonprogram, "Parameter estimation module with external compilers")
else:
    res.append("Tests of parameter estimation module with external compilers: SKIPPED")

if test_pycont_auto:
    print("Testing PyCont continuation with AUTO...\n")
    test(pycont_auto_list, pythonprogram, "PyCont interface to AUTO")
else:
    res.append("Tests of PyCont interface to AUTO: SKIPPED")

if len(failed) == 0:
    print("No test scripts failed")
else:
    print("Test scripts that failed:")
    for fname in failed:
        print("\t%s"%fname)

print("Summary:")
for r in res:
    print(r)


import pytest
import numpy as np
from mesohops.dynamics.hops_aux import AuxiliaryVector
from mesohops.util.exceptions import AuxError

def test_auxvec_ordering():
    """
    This function test whether an incorrectly ordered array_aux_vex properly raises
    the correct AuxError
    """
    aux_1010 = AuxiliaryVector([(0, 1), (2, 1)], 4)
    assert type(aux_1010) == AuxiliaryVector
    with pytest.raises(AuxError) as excinfo:
        aux_1010 = AuxiliaryVector([(2, 1), (0, 1)], 4)
        assert 'array_aux_vec not properly ordered' in str(excinfo.value)

def test_keys():
    """
    This function test whether the correct mode indices (keys) are being grabbed by
    the keys function
    """
    aux_1010 = AuxiliaryVector([(0, 1), (2, 1)], 4)
    keys = aux_1010.keys()
    known_keys = np.array([0, 2])
    assert np.array_equal(keys, known_keys)

def test_values():
    """
    This function test whether the correct values are being grabbed by
    the values function
    """
    aux_1010 = AuxiliaryVector([(0, 1), (2, 1)], 4)
    values = aux_1010.values()
    known_values = np.array([1, 1])
    assert np.array_equal(values, known_values)

def test_compare_if():
    """
    This function test whether _compare is properly comparing auxiliary vectors by
    testing whether one vector is less than another
    """
    aux_1010 = AuxiliaryVector([(0, 1), (2, 1)], 4)
    aux_1000 = AuxiliaryVector([(0, 1)], 4)
    flag = aux_1010._compare(aux_1000, lambda s, o: s > o)
    assert flag == True

def test_compare_else():
    """
    This function test to make sure you cannot compare an Auxiliary Vector to another
    type
    """
    aux_1010 = AuxiliaryVector([(0, 1), (2, 1)], 4)
    known_aux = [(1, 0, 1, 0)]
    flag = aux_1010._compare(known_aux, lambda s, o: s > o)
    assert flag == False

def test_dot():
    """
    This function test to make sure the correct dot product value is given when using
    the dot function
    """
    vector = np.array([1, 1, 1, 1])
    aux_1234 = AuxiliaryVector([(0, 1), (1, 2), (2, 3), (3, 4)], 4)
    known_value = 10
    dot_value = aux_1234.dot(vector)
    assert dot_value == known_value

def test_sum():
    """
    This function test to make sure the values are properly being summed
    """
    aux_array = [(0, 1), (2, 1), (4, 2)]
    n_mod = 6
    aux_101020 = AuxiliaryVector(aux_array, n_mod)
    aux_sum = aux_101020.sum()
    known_sum = 4
    assert aux_sum == known_sum
    # np.sum test
    known_np_sum = 4
    aux_np_sum = np.sum(aux_101020)
    assert aux_np_sum == known_np_sum

def test_todense():
    """
    This function test that a sparse vector is properly being made dense
    """
    aux_101010 = AuxiliaryVector([(0, 1), (2, 1), (4, 1)], 6)
    aux_101010 = aux_101010.todense()
    known_aux = (1, 0, 1, 0, 1, 0)
    assert tuple(aux_101010) == known_aux

def test_toarray():
    """
    This function test that a sparse vector is properly being arranged into an array
    """
    aux_010101 = AuxiliaryVector([(1, 1), (3, 1), (5, 1)], 6)
    aux_010101 = aux_010101.toarray()
    known_array = np.array([[1, 1], [3, 1], [5, 1]])
    assert np.array_equal(aux_010101, known_array)

def test_get_values():
    """
    This function test whether the correct sub-indexed values are being grabbed by
    the get_values function
    """
    aux_101010 = AuxiliaryVector([(0, 1), (2, 1), (4, 1)], 6)
    values = aux_101010.get_values([4, 5])
    known_values = np.array([1, 0])
    assert np.array_equal(values, known_values)

def test_get_values_nonzero():
    """
    This function test whether the correct sub-indexed nonzero values are being grabbed
    by the get_values_nonzero function
    """
    aux_101010 = AuxiliaryVector([(0, 1), (2, 1), (4, 1)], 6)
    values = aux_101010.get_values_nonzero([2, 3, 4, 5])
    known_values = np.array([1, 1])
    assert np.array_equal(values, known_values)

def test_hash_from_estep():
    """
    This function test that the returns the hash of a new Auxiliary Vector
    with the desired step in the given mode is the correct hash
    """
    # Define constants
    aux_2000 = AuxiliaryVector([(0, 2)], 4)
    aux_1001 = AuxiliaryVector([(0, 1), (3, 1)], 4)
    aux_1011 = AuxiliaryVector([(0, 1), (2, 1), (3, 1)], 4)
    aux_1000 = AuxiliaryVector([(0, 1)], 4)
    aux_0000 = AuxiliaryVector([], 4)
    hash_m1 = hash(((0, -1),))

    # test when step = 0
    assert aux_1000.hash == aux_1000.hash_from_e_step(3, 0)
    assert aux_0000.hash == aux_0000.hash_from_e_step(0, 0)
    assert aux_1011.hash == aux_1011.hash_from_e_step(2, 0)

    # test when step = 1
    assert aux_1001.hash == aux_1000.hash_from_e_step(3, 1)
    assert aux_2000.hash == aux_1000.hash_from_e_step(0, 1)
    assert aux_1000.hash == aux_0000.hash_from_e_step(0, 1)
    assert aux_1011.hash == aux_1001.hash_from_e_step(2, 1)

    # test when step = -1
    assert hash_m1 == aux_0000.hash_from_e_step(0, -1)
    assert aux_0000.hash == aux_1000.hash_from_e_step(0, -1)
    assert aux_1000.hash == aux_2000.hash_from_e_step(0, -1)
    assert aux_1000.hash == aux_1001.hash_from_e_step(3, -1)
    assert aux_1001.hash == aux_1011.hash_from_e_step(2, -1)

def test_e_step():
    """
    This function test whether e_step returns a new Auxiliary Vector with the desired
     step in the given mode
    """
    # Define constants
    aux_2000 = AuxiliaryVector([(0, 2)], 4)
    aux_1001 = AuxiliaryVector([(0, 1), (3, 1)], 4)
    aux_1000 = AuxiliaryVector([(0, 1)], 4)
    aux_1011 = AuxiliaryVector([(0, 1), (2, 1), (3, 1)], 4)
    aux_0000 = AuxiliaryVector([], 4)
    Aux_m1 = AuxiliaryVector([(0, -1)], 4)

    # test when step = 0
    assert aux_1000 == aux_1000.e_step(3, 0)
    assert aux_0000 == aux_0000.e_step(0, 0)
    assert aux_1011 == aux_1011.e_step(2, 0)

    # test when step = 1
    assert aux_1001 == aux_1000.e_step(3, 1)
    assert aux_2000 == aux_1000.e_step(0, 1)
    assert aux_1000 == aux_0000.e_step(0, 1)
    assert aux_1011 == aux_1001.e_step(2, 1)

    # test when step = -1
    assert Aux_m1 == aux_0000.e_step(0, -1)
    assert aux_0000 == aux_1000.e_step(0, -1)
    assert aux_1000 == aux_2000.e_step(0, -1)
    assert aux_1000 == aux_1001.e_step(3, -1)
    assert aux_1001 == aux_1011.e_step(2, -1)

def test_index_analytic():
    """
    This function provides a test to ensure an absolute index value is returned
    for an auxiliary vector using an analytic function of the indices
    """
    aux = AuxiliaryVector([(0, 1), (2, 1)], 4)
    # known result based on alpha numerical ordering
    known_ind = 7
    assert aux.absolute_index == known_ind

def test_tuple_from_e_step():
    """
    Test whether tuple_from_e_step returns the sparse correct tuple representation of
    the auxiliary
    """
    # Constants
    aux_0101 = AuxiliaryVector([(1, 1), (3, 1)], 4)
    aux_1010 = AuxiliaryVector([(0, 1), (2, 1)], 4)
    aux_1000 = AuxiliaryVector([(0, 1)], 4)
    aux_0000 = AuxiliaryVector([], 4)
    aux_empty = AuxiliaryVector([], 4)

    # test when step = 0
    known_1000 = ((0, 1),)
    known_0000 = ()
    assert known_1000 == aux_1000.tuple_from_e_step(2, 0)
    assert known_0000 == aux_0000.tuple_from_e_step(0, 0)

    # test when mode + step < 0
    known_tuple = ((0, -1),)
    assert known_tuple == aux_0000.tuple_from_e_step(0, -1)

    # test when len(self.dict_aux_vec) == 0
    known_tuple = ((1, 1),)
    assert known_tuple == aux_empty.tuple_from_e_step(1, 1)

    # test when mode is in array_aux_vec[:, 0] and mode + step = 0
    known_tuple = ((3, 1),)
    assert known_tuple == aux_0101.tuple_from_e_step(1, -1)

    # test when mode is in array_aux_vec[:, 0]
    known_tuple = ((1, 2), (3, 1))
    assert known_tuple == aux_0101.tuple_from_e_step(1, 1)

    # test else
    known_tuple = ((0, 1), (2, 1), (3, 1))
    assert known_tuple == aux_1010.tuple_from_e_step(3, 1)

def test_add_aux_connect():
    """
    Test whether add_aux_connect updates the HopsAux object to contain a pointer to the
    other HopsAux objects it is connected to.
    """
    # Define Constants
    aux_1001 = AuxiliaryVector([(0, 1), (3, 1)], 4)
    aux_1011 = AuxiliaryVector([(0, 1), (2, 1), (3, 1)], 4)
    aux_1000 = AuxiliaryVector([(0, 1)], 4)
    aux_1002 = AuxiliaryVector([(0, 1), (3, 2)], 4)

    # Test when type == 1
    aux_1001.add_aux_connect(2, aux_1011, 1)
    assert aux_1001.dict_aux_p1[2] == aux_1011

    # Test when type == -1
    aux_1001.add_aux_connect(3, aux_1000, -1)
    assert aux_1001.dict_aux_m1[3] == aux_1000

    # # Test when type != +/- 1
    with pytest.raises(AuxError) as excinfo:
        aux_1002.add_aux_connect(3, aux_1000, 2)
        assert 'There is a problem in the hierarchy: add_aux_connect does not support ' \
               'type=2' in str(excinfo.value)

def test_remove_aux_connect():
    """
    Test whether the remove_aux_connect function removes the connection between the
    HopsAux object and another connected with type (+1/-1) along index mode.
    """

    # Define Constants
    aux_1001 = AuxiliaryVector([(0, 1), (3, 1)], 4)
    aux_1011 = AuxiliaryVector([(0, 1), (2, 1), (3, 1)], 4)
    aux_1000 = AuxiliaryVector([(0, 1)], 4)

    # Test when type == 1
    aux_1001.add_aux_connect(2, aux_1011, 1)
    aux_1001.remove_aux_connect(2, 1)
    assert aux_1001.dict_aux_p1 == {}

    # Test when type == -1
    aux_1001.add_aux_connect(3, aux_1000, -1)
    aux_1001.remove_aux_connect(3, -1)
    assert aux_1001.dict_aux_m1 == {}

    # Test when type != +/- 1
    with pytest.raises(AuxError) as excinfo:
        aux_1001.remove_aux_connect(3, 2)
        assert 'There is a problem in the hierarchy: remove_aux_connect does not ' \
               'support ' \
               'type=2' in str(excinfo.value)

def test_remove_pointers():
    """
    This will test if the remove_pointers function removes all pointers targeting the
    current HopsAux object from the set of HopsAux objects it has connections to.
    """

    # Define Constants
    aux_1012 = AuxiliaryVector([(0, 1), (2, 1), (3, 2)], 4)
    aux_1011 = AuxiliaryVector([(0, 1), (2, 1), (3, 1)], 4)
    aux_1010 = AuxiliaryVector([(0, 1), (2, 1)], 4)

    # Test with both +/- 1 additions
    aux_1011.add_aux_connect(3, aux_1012, 1)
    aux_1011.add_aux_connect(3, aux_1010, -1)
    aux_1012.add_aux_connect(3,aux_1011,-1)
    aux_1010.add_aux_connect(3,aux_1011,1)
    aux_1011.remove_pointers()
    assert aux_1011.dict_aux_p1 == {}
    assert aux_1011.dict_aux_m1 == {}
    assert aux_1012.dict_aux_m1 == {}
    assert aux_1010.dict_aux_p1 == {}

def test_difference_by_mode():
    """
    This test will ensure that the difference_by_mode function is returning the
    correct mode in which one HopsAux object differs by another HopsAux object, if the
    difference is only 1 step. This function will also test that an error is called
    if the two objects differ by more than 1 step.
    """

    # Define Constants
    aux_1012 = AuxiliaryVector([(0, 1), (2, 1), (3, 2)], 4)
    aux_3012 = AuxiliaryVector([(0, 3), (2, 1), (3, 2)], 4)
    aux_2012 = AuxiliaryVector([(0, 2), (2, 1), (3, 2)], 4)
    aux_1112 = AuxiliaryVector([(0, 1), (1, 1), (2, 1), (3, 2)], 4)
    aux_1022 = AuxiliaryVector([(0, 1), (2, 2), (3, 2)], 4)
    aux_1013 = AuxiliaryVector([(0, 1), (2, 1), (3, 3)], 4)
    aux_10130 = AuxiliaryVector([(0, 1), (2, 1), (3, 3)], 5)

    # Test mode 0
    difference_0 = aux_1012.difference_by_mode(aux_3012)
    assert difference_0 is False

    # Test mode 1
    difference_1 = aux_1012.difference_by_mode(aux_1112)
    assert difference_1 == [1]

    # Test mode 2
    difference_2 = aux_1012.difference_by_mode(aux_1022)
    assert difference_2 == [2]

    # Test mode 3
    difference_3 = aux_1012.difference_by_mode(aux_1013)
    assert difference_3 == [3]

    # Test when mode of difference is more than one
    difference_many = aux_2012.difference_by_mode(aux_1112)
    assert difference_many is False

    # Test when the two HopsAux objects don't belong to the same hierarchy
    with pytest.raises(AssertionError):
        aux_10130.difference_by_mode(aux_1013)


r"""
Elements of posets, lattices, semilattices, etc.
"""
#*****************************************************************************
#       Copyright (C) 2008 Peter Jipsen <jipsen@chapman.edu>,
#                          Franco Saliola <saliola@gmail.com>
#
#  Distributed under the terms of the GNU General Public License (GPL)
#
#    This code is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
#    General Public License for more details.
#
#  The full text of the GPL is available at:
#
#                  http://www.gnu.org/licenses/
#*****************************************************************************
from sage.structure.element import Element
from sage.structure.element import have_same_parent

class PosetElement(Element):

    def __init__(self, poset, element, vertex):
        r"""
        Establish the parent-child relationship between ``poset``
        and ``element``, where ``element`` is associated to the
        vertex ``vertex`` of the Hasse diagram of the poset.

        INPUT:

        - ``poset`` -- a poset object

        - ``element`` -- any object

        - ``vertex`` -- a vertex of the Hasse diagram of the poset

        TESTS::

            sage: from sage.combinat.posets.elements import PosetElement
            sage: P = Poset([[1,2],[4],[3],[4],[]], facade = False)
            sage: e = P(0)
            sage: e.parent() is P
            True
            sage: TestSuite(e).run()
        """
        Element.__init__(self, poset)
        if isinstance(element, self.parent().element_class):
            self.element = element.element
        else:
            self.element = element
        self.vertex = vertex

    def __hash__(self):
        r"""
        TESTS::

            sage: P = Poset([[1,2],[4],[3],[4],[]], facade = False)
            sage: e = P(0)
            sage: hash(e)
            0
        """
        return hash(self.element)

    def _repr_(self):
        """
        TESTS::

            sage: Poset([[1,2],[4],[3],[4],[]], facade = False)(0)._repr_()
            '0'
        """
        return "%s" % str(self.element)

    def _latex_(self):
        r"""
        Return the latex code of the poset element.

        EXAMPLES::

            sage: m = matrix(2,[1,2,3,4])
            sage: m.set_immutable()
            sage: P = Poset(([m],[]), facade = False)
            sage: [e] = P
            sage: type(e)
            <class 'sage.combinat.posets.posets.FinitePoset_with_category.element_class'>
            sage: latex(e)                 #indirect doctest
            \left(\begin{array}{rr}
            1 & 2 \\
            3 & 4
            \end{array}\right)
        """
        from sage.misc.latex import latex
        return latex(self.element)

    def __eq__(self, other):
        """
        TESTS::

            sage: P = Poset([["a","b"],["d"],["c"],["d"],[]], facade = False)
            sage: Q = Poset([["a","b"],["d"],["c"],[],[]], facade = False)
            sage: P(0).__eq__(P(4))
            False
            sage: from sage.combinat.posets.elements import PosetElement
            sage: PosetElement(P,0,"c") == PosetElement(P,0,"c")
            True
            sage: PosetElement(P,0,"c") == PosetElement(Q,0,"c")
            False
            sage: PosetElement(P,0,"b") == PosetElement(P,0,"c")
            False

        .. warning:: as an optimization, this only compares the parent
           and vertex, using the invariant that, in a proper poset
           element, ``self.element == other.element`` if and only
           ``self.vertex == other.vertex``::

            sage: PosetElement(P,1,"c") == PosetElement(P,0,"c")
            True

        Test that :trac:`12351` is fixed::

            sage: P(0) == int(0)
            False
        """
        # This should instead exploit unique representation, using
        # self is other, or best inherit __eq__ from there. But there
        # are issues around pickling and rich comparison functions.
        return have_same_parent(self, other) \
            and self.vertex == other.vertex

    def __ne__(self, other):
        r"""
        TESTS::

            sage: P = Poset([[1,2],[4],[3],[4],[]])
            sage: P = Poset([["a","b"],["d"],["c"],["d"],[]])
            sage: P(0).__ne__(P(4))
            True
            sage: from sage.combinat.posets.elements import PosetElement
            sage: PosetElement(P,0,"c") != PosetElement(P,0,"c")
            False
            sage: PosetElement(P,0,"b") != PosetElement(P,0,"c")
            True

        For this one, see comment in :meth:`__eq__`::

            sage: PosetElement(P,1,"c") != PosetElement(P,0,"c")
            False
        """
        return not self == other

    def _cmp(self, other):
        """
        TESTS::

            sage: P = Poset([[1,2],[4],[3],[4],[]], facade = False)
            sage: P(0)._cmp(P(4))
            -1
            sage: P(4)._cmp(P(0))
            1
            sage: P(0)._cmp(P(0))
            0
            sage: P(1)._cmp(P(2))

        """
        return self.parent().compare_elements(self, other)

    def __lt__(self, other):
        """
        TESTS

        ::

            sage: dag = DiGraph({0:[2,3], 1:[3,4], 2:[5], 3:[5], 4:[5]})
            sage: P = Poset(dag, facade = False)
            sage: P(0) < P(1)
            False
            sage: P(4) < P(1)
            False
            sage: P(0) < P(0)
            False
        """
        return self._cmp(other) == -1 or False

    def __le__(self, other):
        """
        TESTS

        ::

            sage: dag = DiGraph({0:[2,3], 1:[3,4], 2:[5], 3:[5], 4:[5]})
            sage: P = Poset(dag, facade = False)
            sage: P(1) <= P(0)
            False
            sage: P(0) <= P(1)
            False
            sage: P(0) <= P(3)
            True
            sage: P(0) <= P(0)
            True
        """
        return self == other or self._cmp(other) == -1 or False

    def __gt__(self, other):
        """
        TESTS

        ::

            sage: dag = DiGraph({0:[2,3], 1:[3,4], 2:[5], 3:[5], 4:[5]})
            sage: P = Poset(dag)
            sage: P(0).__gt__(P(5))
            False
            sage: P(5).__gt__(P(0))
            True
            sage: P(0).__gt__(P(0))
            False
        """
        return self._cmp(other) == 1 or False

    def __ge__(self, other):
        """
        TESTS

        ::

            sage: dag = DiGraph({0:[2,3], 1:[3,4], 2:[5], 3:[5], 4:[5]})
            sage: P = Poset(dag)
            sage: P(0).__ge__(P(5))
            False
            sage: P(5).__ge__(P(0))
            True
            sage: P(0).__ge__(P(0))
            True
        """
        return self == other or self._cmp(other) == 1 or False

class MeetSemilatticeElement(PosetElement):
    def __mul__(self, other):
        r"""
        Return the meet of ``self`` and ``other`` in the lattice.

        EXAMPLES::

            sage: D = posets.DiamondPoset(5,facade=False)
            sage: D(1) * D(2)
            0
            sage: D(1) * D(1)
            1
            sage: D(1) * D(0)
            0
            sage: D(1) * D(4)
            1
        """
        return self.parent().meet(self, other)

class JoinSemilatticeElement(PosetElement):
    def __add__(self, other):
        r"""
        Return the join of ``self`` and ``other`` in the lattice.

        EXAMPLES::

            sage: D = posets.DiamondPoset(5,facade=False)
            sage: D(1) + D(2)
            4
            sage: D(1) + D(1)
            1
            sage: D(1) + D(4)
            4
            sage: D(1) + D(0)
            1
        """
        return self.parent().join(self, other)

class LatticePosetElement(MeetSemilatticeElement, JoinSemilatticeElement):
    pass


# SPDX-License-Identifier: Apache-2.0

"""
Intermediate results and investigation
======================================

.. index:: investigate, intermediate results

There are many reasons why a user wants more than using
the converted model into ONNX. Intermediate results may be
needed, the output of every node in the graph. The ONNX
may need to be altered to remove some nodes.
Transfer learning is usually removing the last layers of
a deep neural network. Another reaason is debugging.
It often happens that the runtime fails to compute the predictions
due to a shape mismatch. Then it is useful the get the shape
of every intermediate result. This example looks into two
ways of doing it.

.. contents::
    :local:

Look into pipeline steps
++++++++++++++++++++++++

The first way is a tricky one: it overloads
methods *transform*, *predict* and *predict_proba*
to keep a copy of inputs and outputs. It then goes
through every step of the pipeline. If the pipeline
has *n* steps, it converts the pipeline with step 1,
then the pipeline with steps 1, 2, then 1, 2, 3...
"""
from pyquickhelper.helpgen.graphviz_helper import plot_graphviz
from mlprodict.onnxrt import OnnxInference
import numpy
from onnxruntime import InferenceSession
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.cluster import KMeans
from sklearn.datasets import load_iris
from skl2onnx import to_onnx
from skl2onnx.helpers import collect_intermediate_steps
from skl2onnx.common.data_types import FloatTensorType

###########################
# The pipeline.

data = load_iris()
X = data.data

pipe = Pipeline(steps=[
    ('std', StandardScaler()),
    ('km', KMeans(3))
])
pipe.fit(X)

#################################
# The function goes through every step,
# overloads the methods *transform* and
# returns an ONNX graph for every step.
steps = collect_intermediate_steps(
    pipe, "pipeline",
    [("X", FloatTensorType([None, X.shape[1]]))])

#####################################
# We call method transform to population the
# cache the overloaded methods *transform* keeps.
pipe.transform(X)

#######################################
# We compute every step and compare
# ONNX and scikit-learn outputs.

for step in steps:
    print('----------------------------')
    print(step['model'])
    onnx_step = step['onnx_step']
    sess = InferenceSession(onnx_step.SerializeToString())
    onnx_outputs = sess.run(None, {'X': X.astype(numpy.float32)})
    onnx_output = onnx_outputs[-1]
    skl_outputs = step['model']._debug.outputs['transform']

    # comparison
    diff = numpy.abs(skl_outputs.ravel() - onnx_output.ravel()).max()
    print("difference", diff)

# That was the first way: dynamically overwrite
# every method transform or predict in a scikit-learn
# pipeline to capture the input and output of every step,
# compare them to the output produced by truncated ONNX
# graphs built from the first one.
#
#####################################
# Python runtime to look into every node
# ++++++++++++++++++++++++++++++++++++++
#
# The python runtime may be useful to easily look
# into every node of the ONNX graph.
# This option can be used to check when the computation
# fails due to nan values or a dimension mismatch.

onx = to_onnx(pipe, X[:1].astype(numpy.float32))

oinf = OnnxInference(onx)
oinf.run({'X': X[:2].astype(numpy.float32)},
         verbose=1, fLOG=print)

###################################
# And to get a sense of the intermediate results.

oinf.run({'X': X[:2].astype(numpy.float32)},
         verbose=3, fLOG=print)

# This way is usually better if you need to investigate
# issues within the code of the runtime for an operator.
#
#################################
# Final graph
# +++++++++++

ax = plot_graphviz(oinf.to_dot())
ax.get_xaxis().set_visible(False)
ax.get_yaxis().set_visible(False)


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

# # Homework and bake-off: Relation extraction using distant supervision

# In[ ]:

__author__ = "Bill MacCartney and Christopher Potts"
__version__ = "CS224u, Stanford, Spring 2020"

# ## Contents
#
# 1. [Overview](#Overview)
# 1. [Set-up](#Set-up)
# 1. [Baselines](#Baselines)
#   1. [Hand-build feature functions](#Hand-build-feature-functions)
#   1. [Distributed representations](#Distributed-representations)
# 1. [Homework questions](#Homework-questions)
#   1. [Different model factory [1 points]](#Different-model-factory-[1-points])
#   1. [Directional unigram features [1.5 points]](#Directional-unigram-features-[1.5-points])
#   1. [The part-of-speech tags of the "middle" words [1.5 points]](#The-part-of-speech-tags-of-the-"middle"-words-[1.5-points])
#   1. [Bag of Synsets [2 points]](#Bag-of-Synsets-[2-points])
#   1. [Your original system [3 points]](#Your-original-system-[3-points])
# 1. [Bake-off [1 point]](#Bake-off-[1-point])

# ## Overview
#
# This homework and associated bake-off are devoted to developing really effective relation extraction systems using distant supervision.
#
# As with the previous assignments, this notebook first establishes a baseline system. The initial homework questions ask you to create additional baselines and suggest areas for innovation, and the final homework question asks you to develop an original system for you to enter into the bake-off.

# ## Set-up
#
# See [the first notebook in this unit](rel_ext_01_task.ipynb#Set-up) for set-up instructions.

# In[ ]:

import numpy as np
import os
import rel_ext
from sklearn.linear_model import LogisticRegression
import utils

# As usual, we unite our corpus and KB into a dataset, and create some splits for experimentation:

# In[ ]:

rel_ext_data_home = os.path.join('data', 'rel_ext_data')

# In[ ]:

corpus = rel_ext.Corpus(os.path.join(rel_ext_data_home, 'corpus.tsv.gz'))

# In[ ]:

kb = rel_ext.KB(os.path.join(rel_ext_data_home, 'kb.tsv.gz'))

# In[ ]:

dataset = rel_ext.Dataset(corpus, kb)

# You are not wedded to this set-up for splits. The bake-off will be conducted on a previously unseen test-set, so all of the data in `dataset` is fair game:

# In[ ]:

splits = dataset.build_splits(
    split_names=['tiny', 'train', 'dev'],
    split_fracs=[0.01, 0.79, 0.20],
    seed=1)

# In[ ]:

splits

# ## Baselines

# ### Hand-build feature functions

# In[ ]:

def simple_bag_of_words_featurizer(kbt, corpus, feature_counter):
    for ex in corpus.get_examples_for_entities(kbt.sbj, kbt.obj):
        for word in ex.middle.split(' '):
            feature_counter[word] += 1
    for ex in corpus.get_examples_for_entities(kbt.obj, kbt.sbj):
        for word in ex.middle.split(' '):
            feature_counter[word] += 1
    return feature_counter

# In[ ]:

featurizers = [simple_bag_of_words_featurizer]

# In[ ]:

model_factory = lambda: LogisticRegression(fit_intercept=True, solver='liblinear')

# In[ ]:

baseline_results = rel_ext.experiment(
    splits,
    train_split='tiny',
    test_split='tiny',
    featurizers=featurizers,
    model_factory=model_factory,
    verbose=True)

# Studying model weights might yield insights:

# In[ ]:

rel_ext.examine_model_weights(baseline_results)

# ### Distributed representations
#
# This simple baseline sums the GloVe vector representations for all of the words in the "middle" span and feeds those representations into the standard `LogisticRegression`-based `model_factory`. The crucial parameter that enables this is `vectorize=False`. This essentially says to `rel_ext.experiment` that your featurizer or your model will do the work of turning examples into vectors; in that case, `rel_ext.experiment` just organizes these representations by relation type.

# In[ ]:

GLOVE_HOME = os.path.join('data', 'glove.6B')

# In[ ]:

glove_lookup = utils.glove2dict(
    os.path.join(GLOVE_HOME, 'glove.6B.300d.txt'))

# In[ ]:

def glove_middle_featurizer(kbt, corpus, np_func=np.sum):
    reps = []
    for ex in corpus.get_examples_for_entities(kbt.sbj, kbt.obj):
        for word in ex.middle.split():
            rep = glove_lookup.get(word)
            if rep is not None:
                reps.append(rep)
    # A random representation of the right dimensionality if the
    # example happens not to overlap with GloVe's vocabulary:
    if len(reps) == 0:
        dim = len(next(iter(glove_lookup.values())))
        return utils.randvec(n=dim)
    else:
        return np_func(reps, axis=0)

# In[ ]:
glove_middle_featurizer(kb.kb_triples[0], corpus)

glove_results = rel_ext.experiment(
    splits,
    train_split='tiny',
    test_split='tiny',
    featurizers=[glove_middle_featurizer],
    vectorize=False, # Crucial for this featurizer!
    verbose=True)

# With the same basic code design, one can also use the PyTorch models included in the course repo, or write new ones that are better aligned with the task. For those models, it's likely that the featurizer will just return a list of tokens (or perhaps a list of lists of tokens), and the model will map those into vectors using an embedding.

# ## Homework questions
#
# Please embed your homework responses in this notebook, and do not delete any cells from the notebook. (You are free to add as many cells as you like as part of your responses.)

# ### Different model factory [1 points]
#
# The code in `rel_ext` makes it very easy to experiment with other classifier models: one need only redefine the `model_factory` argument. This question asks you to assess a [Support Vector Classifier](https://scikit-learn.org/stable/modules/generated/sklearn.svm.SVC.html).
#
# __To submit:__ A wrapper function `run_svm_model_factory` that does the following:
#
# 1. Uses `rel_ext.experiment` with the model factory set to one based in an `SVC` with `kernel='linear'` and all other arguments left with default values.
# 1. Trains on the 'train' part of `splits`.
# 1. Assesses on the `dev` part of `splits`.
# 1. Uses `featurizers` as defined above.
# 1. Returns the return value of `rel_ext.experiment` for this set-up.
#
# The function `test_run_svm_model_factory` will check that your function conforms to these general specifications.

# In[ ]:

def run_svm_model_factory():

    ##### YOUR CODE HERE
    from sklearn.svm import SVC
    glove_results = rel_ext.experiment(
        splits,
        train_split='tiny',
        test_split='tiny',
        model_factory=(lambda: SVC(kernel='linear')),
        featurizers=[glove_middle_featurizer],
        vectorize=False,  # Crucial for this featurizer!
        verbose=True)

    return glove_results

# In[ ]:

def test_run_svm_model_factory(run_svm_model_factory):
    results = run_svm_model_factory()
    assert 'featurizers' in results,         "The return value of `run_svm_model_factory` seems not to be correct"
    # Check one of the models to make sure it's an SVC:
    assert 'SVC' in results['models']['adjoins'].__class__.__name__,         "It looks like the model factor wasn't set to use an SVC."

# In[ ]:

if 'IS_GRADESCOPE_ENV' not in os.environ:
    test_run_svm_model_factory(run_svm_model_factory)

# ### Directional unigram features [1.5 points]
#
# The current bag-of-words representation makes no distinction between "forward" and "reverse" examples. But, intuitively, there is big difference between _X and his son Y_ and _Y and his son X_. This question asks you to modify `simple_bag_of_words_featurizer` to capture these differences.
#
# __To submit:__
#
# 1. A feature function `directional_bag_of_words_featurizer` that is just like `simple_bag_of_words_featurizer` except that it distinguishes "forward" and "reverse". To do this, you just need to mark each word feature for whether it is derived from a subject–object example or from an object–subject example.  The included function `test_directional_bag_of_words_featurizer` should help verify that you've done this correctly.
#
# 2. A call to `rel_ext.experiment` with `directional_bag_of_words_featurizer` as the only featurizer. (Aside from this, use all the default values for `rel_ext.experiment` as exemplified above in this notebook.)
#
# 3. `rel_ext.experiment` returns some of the core objects used in the experiment. How many feature names does the `vectorizer` have for the experiment run in the previous step? Include the code needed for getting this value. (Note: we're partly asking you to figure out how to get this value by using the sklearn documentation, so please don't ask how to do it!)

# In[ ]:

def directional_bag_of_words_featurizer(kbt, corpus, feature_counter):
    # Append these to the end of the keys you add/access in
    # `feature_counter` to distinguish the two orders. You'll
    # need to use exactly these strings in order to pass
    # `test_directional_bag_of_words_featurizer`.
    subject_object_suffix = "_SO"
    object_subject_suffix = "_OS"

    ##### YOUR CODE HERE
    for ex in corpus.get_examples_for_entities(kbt.sbj, kbt.obj):
        for word in ex.middle.split(' '):
            feature_counter[word + subject_object_suffix] += 1
    for ex in corpus.get_examples_for_entities(kbt.obj, kbt.sbj):
        for word in ex.middle.split(' '):
            feature_counter[word + object_subject_suffix] += 1

    return feature_counter

# Call to `rel_ext.experiment`:
##### YOUR CODE HERE
featurizers = [directional_bag_of_words_featurizer]
directional_results = rel_ext.experiment(
    splits,
    train_split='tiny',
    test_split='tiny',
    featurizers=featurizers,
    model_factory=model_factory,
    verbose=True)

v = directional_results['vectorizer']
len(v.get_feature_names())  # 1565

# In[ ]:

def test_directional_bag_of_words_featurizer(corpus):
    from collections import defaultdict
    kbt = rel_ext.KBTriple(rel='worked_at', sbj='Randall_Munroe', obj='xkcd')
    feature_counter = defaultdict(int)
    # Make sure `feature_counter` is being updated, not reinitialized:
    feature_counter['is_OS'] += 5
    feature_counter = directional_bag_of_words_featurizer(kbt, corpus, feature_counter)
    expected = defaultdict(
        int, {'is_OS':6,'a_OS':1,'webcomic_OS':1,'created_OS':1,'by_OS':1})
    assert feature_counter == expected,         "Expected:\n{}\nGot:\n{}".format(expected, feature_counter)

# In[ ]:

if 'IS_GRADESCOPE_ENV' not in os.environ:
    test_directional_bag_of_words_featurizer(corpus)

# ### The part-of-speech tags of the "middle" words [1.5 points]
#
# Our corpus distribution contains part-of-speech (POS) tagged versions of the core text spans. Let's begin to explore whether there is information in these sequences, focusing on `middle_POS`.
#
# __To submit:__
#
# 1. A feature function `middle_bigram_pos_tag_featurizer` that is just like `simple_bag_of_words_featurizer` except that it creates a feature for bigram POS sequences. For example, given
#
#   `The/DT dog/N napped/V`
#
#    we obtain the list of bigram POS sequences
#
#    `b = ['<s> DT', 'DT N', 'N V', 'V </s>']`.
#
#    Of course, `middle_bigram_pos_tag_featurizer` should return count dictionaries defined in terms of such bigram POS lists, on the model of `simple_bag_of_words_featurizer`.  Don't forget the start and end tags, to model those environments properly! The included function `test_middle_bigram_pos_tag_featurizer` should help verify that you've done this correctly.
#
# 2. A call to `rel_ext.experiment` with `middle_bigram_pos_tag_featurizer` as the only featurizer. (Aside from this, use all the default values for `rel_ext.experiment` as exemplified above in this notebook.)

# In[ ]:

def middle_bigram_pos_tag_featurizer(kbt, corpus, feature_counter):

    ##### YOUR CODE HERE
    for ex in corpus.get_examples_for_entities(kbt.sbj, kbt.obj):
        for tag_bigram in get_tag_bigrams(ex.middle_POS):
            feature_counter[tag_bigram] += 1

    for ex in corpus.get_examples_for_entities(kbt.obj, kbt.sbj):
        for tag_bigram in get_tag_bigrams(ex.middle_POS):
            feature_counter[tag_bigram] += 1

    return feature_counter

def get_tag_bigrams(s):
    """Suggested helper method for `middle_bigram_pos_tag_featurizer`.
    This should be defined so that it returns a list of str, where each
    element is a POS bigram."""
    # The values of `start_symbol` and `end_symbol` are defined
    # here so that you can use `test_middle_bigram_pos_tag_featurizer`.
    start_symbol = "<s>"
    end_symbol = "</s>"

    ##### YOUR CODE HERE
    tags = [start_symbol] + get_tags(s) + [end_symbol]

    tag_bigrams = []
    for i in range(len(tags) - 1):
        tag_bigrams.append(tags[i] + " " + tags[i + 1])

    return tag_bigrams

def get_tags(s):
    """Given a sequence of word/POS elements (lemmas), this function
    returns a list containing just the POS elements, in order.
    """
    return [parse_lem(lem)[1] for lem in s.strip().split(' ') if lem]

def parse_lem(lem):
    """Helper method for parsing word/POS elements. It just splits
    on the rightmost / and returns (word, POS) as a tuple of str."""
    return lem.strip().rsplit('/', 1)

# Call to `rel_ext.experiment`:
##### YOUR CODE HERE
featurizers = [middle_bigram_pos_tag_featurizer]
middle_bigram_pos_results = rel_ext.experiment(
    splits,
    train_split='tiny',
    test_split='tiny',
    featurizers=featurizers,
    model_factory=model_factory,
    verbose=True)

# In[ ]:

def test_middle_bigram_pos_tag_featurizer(corpus):
    from collections import defaultdict
    kbt = rel_ext.KBTriple(rel='worked_at', sbj='Randall_Munroe', obj='xkcd')
    feature_counter = defaultdict(int)
    # Make sure `feature_counter` is being updated, not reinitialized:
    feature_counter['<s> VBZ'] += 5
    feature_counter = middle_bigram_pos_tag_featurizer(kbt, corpus, feature_counter)
    expected = defaultdict(
        int, {'<s> VBZ':6,'VBZ DT':1,'DT JJ':1,'JJ VBN':1,'VBN IN':1,'IN </s>':1})
    assert feature_counter == expected,         "Expected:\n{}\nGot:\n{}".format(expected, feature_counter)

# In[ ]:

if 'IS_GRADESCOPE_ENV' not in os.environ:
    test_middle_bigram_pos_tag_featurizer(corpus)

# ### Bag of Synsets [2 points]
#
# The following allows you to use NLTK's WordNet API to get the synsets compatible with _dog_ as used as a noun:
#
# ```
# from nltk.corpus import wordnet as wn
# dog = wn.synsets('dog', pos='n')
# dog
# [Synset('dog.n.01'),
#  Synset('frump.n.01'),
#  Synset('dog.n.03'),
#  Synset('cad.n.01'),
#  Synset('frank.n.02'),
#  Synset('pawl.n.01'),
#  Synset('andiron.n.01')]
# ```
#
# This question asks you to create synset-based features from the word/tag pairs in `middle_POS`.
#
# __To submit:__
#
# 1. A feature function `synset_featurizer` that is just like `simple_bag_of_words_featurizer` except that it returns a list of synsets derived from `middle_POS`. Stringify these objects with `str` so that they can be `dict` keys. Use `convert_tag` (included below) to convert tags to `pos` arguments usable by `wn.synsets`. The included function `test_synset_featurizer` should help verify that you've done this correctly.
#
# 2. A call to `rel_ext.experiment` with `synset_featurizer` as the only featurizer. (Aside from this, use all the default values for `rel_ext.experiment`.)

# In[ ]:

from nltk.corpus import wordnet as wn

def synset_featurizer(kbt, corpus, feature_counter):

    ##### YOUR CODE HERE
    for ex in corpus.get_examples_for_entities(kbt.sbj, kbt.obj):
        for synset in get_synsets(ex.middle_POS):
            feature_counter[str(synset)] += 1

    for ex in corpus.get_examples_for_entities(kbt.obj, kbt.sbj):
        for synset in get_synsets(ex.middle_POS):
            feature_counter[str(synset)] += 1

    return feature_counter

def get_synsets(s):
    """Suggested helper method for `synset_featurizer`. This should
    be completed so that it returns a list of stringified Synsets
    associated with elements of `s`.
    """
    # Use `parse_lem` from the previous question to get a list of
    # (word, POS) pairs. Remember to convert the POS strings.
    wt = [parse_lem(lem) for lem in s.strip().split(' ') if lem]

    ##### YOUR CODE HERE
    synsets = []
    for word_tag in wt:
        synsets.extend(wn.synsets(word_tag[0], convert_tag(word_tag[1])))

    return synsets

def convert_tag(t):
    """Converts tags so that they can be used by WordNet:

    | Tag begins with | WordNet tag |
    |-----------------|-------------|
    | `N`             | `n`         |
    | `V`             | `v`         |
    | `J`             | `a`         |
    | `R`             | `r`         |
    | Otherwise       | `None`      |
    """
    if t[0].lower() in {'n', 'v', 'r'}:
        return t[0].lower()
    elif t[0].lower() == 'j':
        return 'a'
    else:
        return None

# Call to `rel_ext.experiment`:
##### YOUR CODE HERE
featurizers = [synset_featurizer]
synset_results = rel_ext.experiment(
    splits,
    train_split='tiny',
    test_split='tiny',
    featurizers=featurizers,
    model_factory=model_factory,
    verbose=True)

# In[ ]:

def test_synset_featurizer(corpus):
    from collections import defaultdict
    kbt = rel_ext.KBTriple(rel='worked_at', sbj='Randall_Munroe', obj='xkcd')
    feature_counter = defaultdict(int)
    # Make sure `feature_counter` is being updated, not reinitialized:
    feature_counter["Synset('be.v.01')"] += 5
    feature_counter = synset_featurizer(kbt, corpus, feature_counter)
    # The full return values for this tend to be long, so we just
    # test a few examples to avoid cluttering up this notebook.
    test_cases = {
        "Synset('be.v.01')": 6,
        "Synset('embody.v.02')": 1
    }
    for ss, expected in test_cases.items():
        result = feature_counter[ss]
        assert result == expected,             "Incorrect count for {}: Expected {}; Got {}".format(ss, expected, result)

# In[ ]:

if 'IS_GRADESCOPE_ENV' not in os.environ:
    test_synset_featurizer(corpus)

# ### Your original system [3 points]
#
# There are many options, and this could easily grow into a project. Here are a few ideas:
#
# - Try out different classifier models, from `sklearn` and elsewhere.
# - Add a feature that indicates the length of the middle.
# - Augment the bag-of-words representation to include bigrams or trigrams (not just unigrams).
# - Introduce features based on the entity mentions themselves. <!-- \[SPOILER: it helps a lot, maybe 4% in F-score. And combines nicely with the directional features.\] -->
# - Experiment with features based on the context outside (rather than between) the two entity mentions — that is, the words before the first mention, or after the second.
# - Try adding features which capture syntactic information, such as the dependency-path features used by Mintz et al. 2009. The [NLTK](https://www.nltk.org/) toolkit contains a variety of [parsing algorithms](http://www.nltk.org/api/nltk.parse.html) that may help.
# - The bag-of-words representation does not permit generalization across word categories such as names of people, places, or companies. Can we do better using word embeddings such as [GloVe](https://nlp.stanford.edu/projects/glove/)?
#
# In the cell below, please provide a brief technical description of your original system, so that the teaching team can gain an understanding of what it does. This will help us to understand your code and analyze all the submissions to identify patterns and strategies.

# In[ ]:

# Enter your system description in this cell.
# Please do not remove this comment.

# ## Bake-off [1 point]
#
# For the bake-off, we will release a test set. The announcement will go out on the discussion forum. You will evaluate your custom model from the previous question on these new datasets using the function `rel_ext.bake_off_experiment`. Rules:
#
# 1. Only one evaluation is permitted.
# 1. No additional system tuning is permitted once the bake-off has started.
#
# The cells below this one constitute your bake-off entry.
#
# People who enter will receive the additional homework point, and people whose systems achieve the top score will receive an additional 0.5 points. We will test the top-performing systems ourselves, and only systems for which we can reproduce the reported results will win the extra 0.5 points.
#
# Late entries will be accepted, but they cannot earn the extra 0.5 points. Similarly, you cannot win the bake-off unless your homework is submitted on time.
#
# The announcement will include the details on where to submit your entry.

# In[ ]:

# Enter your bake-off assessment code in this cell.
# Please do not remove this comment.
if 'IS_GRADESCOPE_ENV' not in os.environ:
    pass
    # Please enter your code in the scope of the above conditional.
    ##### YOUR CODE HERE

# In[ ]:

# On an otherwise blank line in this cell, please enter
# your macro-average f-score (an F_0.5 score) as reported
# by the code above. Please enter only a number between
# 0 and 1 inclusive. Please do not remove this comment.
if 'IS_GRADESCOPE_ENV' not in os.environ:
    pass
    # Please enter your score in the scope of the above conditional.
    ##### YOUR CODE HERE


"""
        Assignment-2
        Name - Gajraj Singh Chouhan
        Roll No - B19130
        Branch - DSE
        Mobile No - +91-9351159849
"""

import pandas as pd
import numpy as np
from matplotlib import pyplot as plt
from functions import properties, plot_rmse

# imports

data = pd.read_csv('pima_indians_diabetes_miss.csv')
original_data = pd.read_csv('pima_indians_diabetes_original.csv')

# Question 1

missing_val = data.isna().sum()
# We can count the missing values by using isna(), adding all the cells in a column which would a NaN value (isna() would give each cell 1 or 0 whether it is or is-not a NaN value).
missing_val.plot.bar(rot=0)
# plotting the missing values.
# adding the labels and title.
plt.title("Missing Values in the Attributes")
plt.xlabel("Attributes")
plt.ylabel("Missing Values")
plt.savefig("Q1-Barplot.png")
plt.close()

# Question 2a
print('Question 2a', end='\n\n')

col = data.shape[1] # number of columns
to_drop = 3 # inverse of 1/3rd rows to drop
filtered = data.isna().sum(axis=1)
# finding how many NaN values are in each row by summing the boolean values for each row. (using axis=1 will do sum of each row)
# Now we have a dataframe consisting of number of NaN values in that row

filtered = filtered[(filtered * to_drop) >= col]
# We can filter the rows to drop, by using the condition that number of NaN values must be greater or equal to 1/3rd of no of columns.
# This will give us True or False Boolean Value if that row is to be dropped or not.

print(f"Total number of tuples deleted = {filtered.size}")

filtered = filtered.index
data = data.drop(filtered, axis=0)
# Now we can drop the rows in the original DataFrame using their indexes we got in the previous step.

print(f"Indices of rows dropped (0-based): ")
print(filtered.to_list());print()

# Question 2b
print('Question 2b', end='\n\n')

filtered = data[data['class'].isna()]
# rows whose 'class' attribute is NaN using conditional filtering
filtered = filtered.index
# getting index of rows which we will delete

print(f"Total number of tuples deleted = {filtered.size}")

data = data.drop(filtered, axis=0) # remove the rows

print(f"Indices of rows dropped (0-based): ")
print(filtered.to_list());print()

# Question 3
print('Question 3', end='\n\n')

count_of_NaN = data.isna().sum()
# again we can use isna() for number of null values.

print("Number of missing values in each attributes: ")
print(count_of_NaN.to_frame().T)
# printing number of missing value per column
# ignore the first 0, it just the first row number as I am displaying the Series in horizontal way.
print(f"Total number of missing values in the dataFrame : {count_of_NaN.sum()}");print()

# Question 4a
print('Question 4a', end='\n\n')
column_mean = data.mean()
data_withmean = data.fillna(column_mean)
# fillna() will fill the NaN values for each attribute by the mean of that column

print('RMSE of Original and Mean Filled data')
plot_rmse(data_withmean, original_data, np.where(pd.isnull(data)), count_of_NaN);print()

# plotting the rmse after replacing the nan values.
plt.title("RMSE of original and mean filled data")
plt.savefig("Q4-MeanFilled.png")
plt.close()

# Question 4b
print('Question 4b', end='\n\n')

data_with_interpolate = data.interpolate()
# by default each column will be filled with data interpolated with linear method.

print('RMSE of Original and interpolated data')
plot_rmse(data_with_interpolate, original_data, np.where(pd.isnull(data)), count_of_NaN);print()

plt.title("RMSE of original and interpolated data")
plt.savefig("Q4-Interpolated.png")
plt.close()

# Question 4
print('Question 4 (comparing the filled data and original) : ', end='\n\n')

# Comparing the properties (mean, median, mode, standard-dev...) of the
# original data and after we filled the values in 4a and 4b.
print('Original Data : ')
print(properties(original_data).to_string())
print('\nAfter filling data with mean : ')
print(properties(data_withmean).to_string())
print('\nAfter interpolating the data : ')
print(properties(data_with_interpolate).to_string());print()

# Question 5
print('Question 5 (Outliers)', end='\n\n')

data_with_interpolate = data.interpolate()
columns = ["Age", "BMI"]
quartiles = data_with_interpolate.quantile([0.25, 0.5, 0.75]).loc[:, columns]
# quantile() will give the required quartiles (0.25, 0.5, 0.75) for required columns
data.loc[:, columns].plot.box()
plt.title("Box-Plot of data before replacing outliers with median")
plt.savefig("Q5-before-BoxPlot.png")
plt.close()

for col in columns:
    q1, median, q3 = quartiles.loc[:, col].to_list()
    # quartiles for one column
    print(f"Column : {col}")
    check_outlier = lambda num : (q1 - (1.5 * (q3 - q1))) < num < (q3 + (1.5 * (q3 - q1)))
    # this will return boolean value whether a number ("num") would be outlier according to the quartiles.
    for ind, x in enumerate(data_with_interpolate.loc[:, col]):
        # iterating over the column and checking if it's a outlier.
        if not check_outlier(x):
            # if it is a outlier
            data.loc[data_with_interpolate.index[ind], col] = median
            # we can change value of a cell to the median of that column from the column name and row index.
            # data_with_interpolate.index[ind] - row index
            print(x, end=' ')
    print()

data.loc[:, columns].plot.box()
plt.title("Box-Plot of data after replacing outliers with median")
plt.savefig("Q5-after-BoxPlot.png")
plt.close()

print('properties after replacing with median : ')
print(properties(data.loc[:, columns]));print()


"""
Module file containing functions that allow to reproduce FIG. 1 of the article

    Interaction of "Solitons" in a Collisionless Plasma and the Recurrence of Initial States,
    N. J. Zabusky and M. D. Kruskal,
    Phys. Rev. Lett. 15, 240 (1965)

This module was prepared as a part of the scientific short course

  A brief guide to publication-ready scientific figures using Python's matplotlib

held during the 2020 seminar week of the Ultrafast Laser Laboratory
at Institute of Quantum Optics at Leibniz University Hannover.

Author: O. Melchert
Date: 2020-09-08
"""
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec

__author__ = 'Oliver Melchert'
__date__ = '2020-09-09'

def fetch_data(path):
    """fetch data

    Reads in data from file in numpy npz-format

    Args:
      path (str): path to npz-file

    Returns: (x, t, uxt)
      x (1D array): x samples
      t (1D array): t samples
      uxt (2D array): wave profile u(x,t)
    """
    dat = np.load(path)
    return dat['x'], dat['t'], dat['uxt']

def set_style():
    """set figure style

    Function that customizes the default style to be conform with the Physical
    Review style and notation guide [1]. For instructions on how to set the
    default style using style sheets see [2].

    Notes:
    - main font size is chosen as 8pt, matching the fontsize of figure captions
    - fontsize of legends and auxiliary text labels are set to 6pt, exceeding
      the minimally required pointsize of 4.25 pt. (1.5 mm)
    - default rc (rc = "run commands", i.e. startup information) settings are
      changed dynamically
    - the custom font-scheme 'type2' depends on your latex installation and
      is not guaranteed to run on your specific system

    Refs:
      [1] https://journals.aps.org/prl/authors
      [2] https://matplotlib.org/3.3.1/tutorials/introductory/customizing.html
    """

    fig_width_1col = 3.4        # figure width in inch
    fig_width_2col = 7.0        # figure width in inch
    fig_aspect_ratio = 0.66     # width to height aspect ratio
    font_size = 8               # font size in pt
    font_size_small = 6         # font size in pt
    font_scheme = 'type2'       # options:
                                #   None    - default matplotlib fonts
                                #   'type1' - text: Helvetica, math: Computer modern
                                #   'type2' - text: Helvetica, math: Helvetica

    mpl.rcParams['figure.figsize'] = fig_width_1col, fig_aspect_ratio*fig_width_1col
    mpl.rcParams['axes.labelsize'] = font_size
    mpl.rcParams['font.size'] = font_size
    mpl.rcParams['legend.fontsize'] = font_size_small
    mpl.rcParams['xtick.labelsize'] = font_size
    mpl.rcParams['ytick.labelsize'] = font_size
    mpl.rcParams['xtick.direction'] = 'out'
    mpl.rcParams['ytick.direction'] = 'out'
    mpl.rcParams['lines.linewidth'] = 1.0
    mpl.rcParams['axes.linewidth'] =  0.5

    if font_scheme == 'type1':
        mpl.rcParams['text.usetex'] = True
        mpl.rcParams['font.family'] = 'sans-serif'
        mpl.rcParams['font.sans-serif'] = 'Helvetica'
        mpl.rcParams['mathtext.fontset'] = 'cm'

    if font_scheme == 'type2':
        mpl.rcParams['text.usetex'] = True
        mpl.rcParams['text.latex.preamble'] = [
           r'\usepackage{siunitx}',
           r'\sisetup{detect-all}',
           r'\usepackage{helvet}',
           r'\usepackage{sansmath}',
           r'\sansmath'
        ]

def set_circle(ax, x0, y0, label):
    """set circle

    Function that generates a circle with text-label at its center.
    For more options on scatter plots, see [1], for more options on
    setting text, see [2]

    Refs:
      [1] https://matplotlib.org/3.3.1/api/_as_gen/matplotlib.pyplot.scatter.html
      [2] https://matplotlib.org/3.1.1/api/_as_gen/matplotlib.pyplot.text.html

    Args:
      ax (object): figure part for which the labeled circle is intended
      x0 (float): x-position for center of circle
      y0 (float): y-position for center of circle
      label (str): text that should be displayed within circle
    """
    # -- scatter plot consisting of a single object
    ax.scatter(x0, y0, s=50, linewidth=0.75, facecolors='none', edgecolor='black', zorder=10)
    # -- place text at the center of the scatter plot object
    ax.text(x0, y0, label, backgroundcolor='none', ha='center', va='center', color='black', zorder=10, fontsize=6)

def set_legend(ax, lines):
    """set legend

    Function generating a custom legend, see [1] for more options

    Refs:
      [1] https://matplotlib.org/3.1.1/api/_as_gen/matplotlib.pyplot.legend.html

    Args:
      ax (object): figure part for which the legend is intended
      lines (list): list of  Line2D objects
    """
    # -- extract labels from lines
    labels = [x.get_label() for x in lines]
    # -- customize legend
    ax.legend(lines,                # list of Line2D objects
              labels,               # labels
              title = '',           # title shown on top of legend
              loc = 0,              # location of the legend
              ncol = 1,             # number of columns
              labelspacing = 0.3,   # vertical space between handles in font-size units
              borderpad = 0.3,      # distance to legend border in font-size units
              handletextpad = 0.6,  # distance between handle and label in font-size units
              handlelength = 2.0,   # length of handle in font-size units
              frameon = False       # remove background patch
              )

def set_grid(ax):
    """set grid

    Function generating a custom grid, see [1] for more options

    Refs:
      [1] https://matplotlib.org/3.1.1/api/_as_gen/matplotlib.pyplot.grid.html

    Args:
      ax (object): figure part for which the grid is intended
    """
    #ax.set_facecolor('lightgray')
    ax.grid(which = 'major',        # grid lines at major ticks only
            linestyle = '-',        # solid line
            linewidth = 0.5,        # width of gridlines
            color = 'lightgray',    # color of gridlines
            zorder = 0
            )

def save_figure(fig_format = None, fig_name = 'test'):
    """ save figure

    Function that saves figure or shows interactive plot

    Note:
    - if no valid option is provided, an interactive plot is shown

    Args:
      fig_format (str): format to save figure in (options: png, pdf, svg)
      fig_name (str): name for figure (default: 'test')
    """
    if fig_format == 'png':
        plt.savefig(fig_name+'.png', format='png', dpi=600)
    elif fig_format == 'pdf':
        plt.savefig(fig_name+'.pdf', format='pdf', dpi=600)
    elif fig_format == 'svg':
        plt.savefig(fig_name+'.svg', format='svg')
    else:
        plt.show()

def generate_figure(x, t, uxt, fig_format=None, fig_name='fig01'):
    """generate figure

    Function generating a figure reproducing FIG. 1 of [1].

    Refs:
      [1] Interaction of "Solitons" in a Collisionless Plasma and the Recurrence of Initial States,
          N. J. Zabusky and M. D. Kruskal,
          Phys. Rev. Lett. 15, 240 (1965)

    Args:
      x (1D array): x samples
      t (1D array): t samples
      uxt (2D array): wave profile u(x,t)
      fig_format (str): format for output figure
                        (choices: png, pdf, svg; default: interactive figure)
      fig_name (str): name for output figure wihtout suffix (default='fig01')
    """

    # (1) SET A STYLE THAT FITS THE TARGET JOURNAL
    set_style()

    # (2) SET FIGURE LAYOUT
    fig = plt.figure()
    plt.subplots_adjust(left = 0.12,    # pos. of subplots left border
                        bottom = 0.16,  # pos. of subplots bottom border
                        right = 0.97,   # pos. of subplots right border
                        top = 0.98      # pos. of subplots top border
                        )
    gs00 = GridSpec(nrows = 1, ncols = 1)   # set geometry of subplot grid
    ax01 = fig.add_subplot(gs00[:,:])       # add subplot

    # (3) SET AXES CONTENTS
    # -- custom constants and local functions
    tB = 1./np.pi                                   # breakdown time
    _ux = lambda t0: uxt[np.argmin(np.abs(t-t0))]   # wave form for array index closest to t0

    # -- plot real-vaued field at selected times
    l1 = ax01.plot(x, _ux(0.0),    color='blue',  dashes=[1,1], zorder=100, label=r'$t=0$')
    l2 = ax01.plot(x, _ux(1.0*tB), color='green', dashes=[3,1], zorder=101, label=r'$t=t_{\mathrm{B}}$')
    l3 = ax01.plot(x, _ux(3.6*tB), color='black', zorder=102, label=r'$t=3.6\,t_{\mathrm{B}}$')

    # -- add legend
    set_legend(ax01, l1 + l2 + l3)

    # -- set data curve labels
    ax01.text(1.27, -0.9,               # x,y pos. in data coordinates
                r'A',                   # text label
                color='blue',           # text color
                backgroundcolor='none', # backgroundcolor
                ha='left',              # horizontal alignment
                va='bottom',            # vertical alignment
                zorder=10               # layering order in figure
                )
    ax01.text(0.60, -0.9, r'B', color='green', backgroundcolor='none',
                    ha='left', va='bottom', zorder=10)
    ax01.text(1.75, -0.9, r'C', color='black', backgroundcolor='none',
                    ha='left', va='bottom', zorder=10)

    # -- set auxiliary dash-dotted lines
    _f = lambda x: -1.5+1.5*np.where(x>0.8, x, x+2.)
    ax01.plot(x[x<0.75], _f(x[x<0.75]), color='black',
                    dashes = [15,1,2,1,2,1], linewidth=0.75)
    ax01.plot(x[x>0.80], _f(x[x>0.80]), color='black',
                    dashes = [15,1,2,1,2,1], linewidth=0.75)

    # -- set circles with labels
    for  idx, x0 in enumerate([0.643, 0.362, 0.098, 0.920, 1.14, 1.37, 1.6, 1.845]):
        set_circle(ax01, x0, _f(x0) + 0.2, '$%d$'%(idx+1))

    # (4) SET AXIS DETAILS
    # -- customize x-axis
    x_lim = (0, 2)
    x_ticks = (0., 0.5, 1.0, 1.5, 2.0)
    # -- major ticks
    ax01.tick_params(
        axis = 'x',         # axis for which to set ticks
        direction = 'out',  # place ticks outside
        length = 3.5,       # tick length in pts.
        pad = 2,            # tick to label distance in pts.
        top=False           # no ticks on top spine
        )
    ax01.set_xlim(x_lim)
    ax01.set_xticks(x_ticks)
    # -- minor ticks
    ax01.tick_params(
        axis = 'x',         # axis for which to set ticks
        which = 'minor',    # address minor ticks
        length = 2.         # tick length in pts.
        )
    ax01.set_xticks(
        np.linspace(x_lim[0],x_lim[1], int((x_lim[1]-x_lim[0])/0.1), endpoint=False),
        minor=True
        )
    # -- label
    ax01.set_xlabel(r'Normalized distance $x$')

    # -- customize y-axis
    y_lim = (-1., 3.)
    y_ticks = (-1., 0., 1., 2., 3.)
    # -- major ticks
    ax01.tick_params(axis='y', direction='out', length=3.5, pad=2, right=False)
    ax01.set_ylim(y_lim)
    ax01.set_yticks(y_ticks)
    # -- minor ticks
    ax01.tick_params(axis='y', which='minor', length=2.)
    ax01.set_yticks(np.linspace(y_lim[0], y_lim[1], int((y_lim[1]-y_lim[0])/0.2), endpoint=False), minor=True)
    # -- label
    ax01.set_ylabel(r'Real-valued field $u(x,t)$')

    # (5) SET BACKGROUND GRID
    set_grid(ax01)

    # (6) SAVE FIGURE
    save_figure(fig_format, fig_name)

def main():
    x, t, uxt = fetch_data('KdV_raw_data.npz')
    generate_figure(x, t, uxt, fig_format='png', fig_name='fig01')

if __name__ == '__main__':
    main()


"""
.. _ref_basic-geometry-areas:

Areas
-----
This example shows how you can use PyMAPDL to create basic geometry
using Pythonic PREP7 area commands.

"""

import numpy as np
from ansys.mapdl.core import launch_mapdl

# start MAPDL and enter the pre-processing routine
mapdl = launch_mapdl()
mapdl.clear()
mapdl.prep7()
print(mapdl)

###############################################################################
# APDL Command: A
# ~~~~~~~~~~~~~~~
# Create a simple triangle in the XY plane using three keypoints.

k0 = mapdl.k("", 0, 0, 0)
k1 = mapdl.k("", 1, 0, 0)
k2 = mapdl.k("", 0, 1, 0)
a0 = mapdl.a(k0, k1, k2)
mapdl.aplot(show_lines=True, line_width=5, show_bounds=True, cpos="xy")

###############################################################################
# APDL Command: AL
# ~~~~~~~~~~~~~~~~
# Create an area from four lines.
mapdl.clear()
mapdl.prep7()

k0 = mapdl.k("", 0, 0, 0)
k1 = mapdl.k("", 1, 0, 0)
k2 = mapdl.k("", 1, 1, 0)
k3 = mapdl.k("", 0, 1, 0)
l0 = mapdl.l(k0, k1)
l1 = mapdl.l(k1, k2)
l2 = mapdl.l(k2, k3)
l3 = mapdl.l(k3, k0)
anum = mapdl.al(l0, l1, l2, l3)
mapdl.aplot(show_lines=True, line_width=5, show_bounds=True, cpos="xy")

###############################################################################
# APDL Command: ADRAG
# ~~~~~~~~~~~~~~~~~~~
# Generate areas by dragging a line pattern along a path.
#
# Drag a circle between two keypoints to create an area
mapdl.clear()
mapdl.prep7()

k0 = mapdl.k("", 0, 0, 0)
k1 = mapdl.k("", 0, 0, 1)
carc = mapdl.circle(k0, 1, k1, arc=90)
l0 = mapdl.l(k0, k1)
mapdl.adrag(carc[0], nlp1=l0)
mapdl.aplot(show_lines=True, line_width=5, show_bounds=True, smooth_shading=True)

###############################################################################
# APDL Command: ASBA
# ~~~~~~~~~~~~~~~~~~
# Subtract a ``0.5 x 0.5`` rectangle from a ``1 x 1`` rectangle.
mapdl.clear()
mapdl.prep7()

anum0 = mapdl.blc4(0, 0, 1, 1)
anum1 = mapdl.blc4(0.25, 0.25, 0.5, 0.5)
aout = mapdl.asba(anum0, anum1)
mapdl.aplot(show_lines=True, line_width=5, show_bounds=True, cpos="xy")

###############################################################################
# Area IDs
# ~~~~~~~~
# Return an array of the area IDs
anum = mapdl.geometry.anum
anum

###############################################################################
# Area Geometry
# ~~~~~~~~~~~~~
# Get the VTK ``PolyData`` containing lines.  This VTK mesh can be
# saved or plotted.  For more details, visit https://docs.pyvista.com
#
# Note that this is a method so you can select the quality of the
# areas (mesh density), and if you would like a merged output or
# individual meshes.
areas = mapdl.geometry.areas(quality=3)
areas

###############################################################################
# Merged Area Geometry
# ~~~~~~~~~~~~~~~~~~~~
area = mapdl.geometry.areas(quality=3, merge=True)
area

# optionally save the area, or plot it
# area.save('mesh.vtk')
# area.plot()

###############################################################################
# Area Selection
# ~~~~~~~~~~~~~~
# There are two approaches for selecting areas, the old "legacy" style
# and the new style.  The old style is valuable for those who are
# comfortable with the existing MAPDL commands, and new style is
# useful for selecting areas in a pythonic manner.
#
# This example generates a series of random squares and selects them
mapdl.clear()
mapdl.prep7()

def generate_random_area():
    start_x, start_y, height, width = np.random.random(4)
    mapdl.blc4(start_x * 10, start_y * 10, height, width)

# create 20 random rectangles
for i in range(20):
    generate_random_area()

# Print the area numbers
print(mapdl.geometry.anum)

###############################################################################
# Select every other area with the old style command.
mapdl.asel("S", "AREA", "", 1, 20, 2)
print(mapdl.geometry.anum)

###############################################################################
# Select every other area with the new style command.
#
# Note that the Area IDs are 1 based in MAPDL, while Python ranges are 0 based.
mapdl.geometry.area_select(range(1, 21, 2))
print(mapdl.geometry.anum)

###############################################################################
# Select areas from a list
#
# Note that you can ``return_selected`` if you want to see what you
# have selected.  This is helpful when reselecting from existing
# areas.
items = mapdl.geometry.area_select([1, 5, 10, 20], return_selected=True)
print(items)

###############################################################################
# APDL Command: APLOT
# ~~~~~~~~~~~~~~~~~~~
# This method uses VTK and pyvista to generate a dynamic 3D plot.
#
# There are a variety of plotting options available for all the common
# plotting methods.  Here, we enable the bounds and show the lines of
# the plot while increasing the plot quality with the ``quality``
# parameter.
#
# Note that the `cpos` keyword argument can be used to describe the
# camera direction from the following:
#
# - 'iso' - Isometric view
# - 'xy' - XY Plane view
# - 'xz' - XZ Plane view
# - 'yx' - YX Plane view
# - 'yz' - YZ Plane view
# - 'zx' - ZX Plane view
# - 'zy' - ZY Plane view

mapdl.aplot(quality=1, show_bounds=True, cpos="iso", show_lines=True)


import os
import numpy as np
from models.LinearRegression import LinearRegression
from models.LogisticRegression import LogisticRegression
from optim.Optimizer import SGD, Momentum, RMSProp

# =================================================================== #
#                   DO NOT EDIT ANY SINGLE CHARACTER                  #
# =================================================================== #

def load_reg_data(path, filename, target_at_front, normalize=False, shuffle=False):
    fullpath = os.path.join(path, filename)

    with open(fullpath, 'r') as f:
        lines = f.readlines()
    lines = [s.strip().split(',') for s in lines]

    header = lines[0]
    data = lines[1:]

    data = np.array([[float(f) for f in d] for d in data], dtype=np.float32)
    if target_at_front:
        x, y = data[:, 1:], data[:, 0]
    else:
        x, y = data[:, :-1], data[:, -1]

    num_data = x.shape[0]
    if normalize:
        mins = np.expand_dims(np.min(x, axis=0), 0).repeat(num_data, 0)
        maxs = np.expand_dims(np.max(x, axis=0), 0).repeat(num_data, 0)
        x = (x - mins) / maxs

    # Add 1 column for bias
    bias = np.ones((x.shape[0], 1), dtype=np.float32)
    x = np.concatenate((bias, x), axis=1)

    if shuffle:
        perm = np.random.permutation(num_data)
        x = x[perm]
        y = y[perm]

    return x, y

def load_class_data(path, filename, target_at_front, to_binary=False, normalize=False, excludes=None):
    if excludes is None:
        excludes = []

    fullpath = os.path.join(path, filename)

    with open(fullpath, 'r') as f:
        lines = f.readlines()
    lines = [s.strip().split(',') for s in lines]

    header = lines[0]
    raw_data = lines[1:]
    num_feat = len(raw_data[0])
    feat_to_idx = [{} for _ in range(num_feat)]

    data = []
    for d in raw_data:
        line = []
        for i, f in enumerate(d):
            if i in excludes:
                continue

            try:
                line.append(float(f))
            except:
                if f in feat_to_idx[i]:
                    f_idx = feat_to_idx[i][f]
                else:
                    f_idx = len(feat_to_idx[i])
                    feat_to_idx[i][f] = f_idx
                line.append(f_idx)
        data.append(line)

    data = np.array(data, dtype=np.float32)
    if target_at_front:
        x, y = data[:, 1:], data[:, 0].astype(np.int32)
    else:
        x, y = data[:, :-1], data[:, -1].astype(np.int32)

    num_data = x.shape[0]
    if normalize:
        mins = np.expand_dims(np.min(x, axis=0), 0).repeat(num_data, 0)
        maxs = np.expand_dims(np.max(x, axis=0), 0).repeat(num_data, 0)
        x = (x - mins) / maxs

    # Add 1 column for bias
    bias = np.ones((x.shape[0], 1), dtype=np.float32)
    x = np.concatenate((bias, x), axis=1)

    if to_binary:
        y[y > 1] = 1

    return x, y

def ConcreteData(path, filename):
    return load_reg_data(path, filename, target_at_front=False, normalize=True)

def GraduateData(path, filename):
    return load_reg_data(path, filename, target_at_front=False, normalize=True)

def DigitData(path, filename):
    x, y = load_class_data(path, filename, target_at_front=True, excludes=[1, 2, 3, 4, 5, 7, 8, 9], normalize=False, to_binary=True)
    x = x / 256
    return (x, y)

def TitanicData(path, filename):
    return load_class_data(path, filename, target_at_front=True, excludes=[2])

def RMSE(h, y):
    if len(h.shape) > 1:
        h = h.squeeze()
    se = np.square(h - y)
    mse = np.mean(se)
    rmse = np.sqrt(mse)
    return rmse

def Accuracy(h, y):
    if len(y.shape) == 1:
        y = np.expand_dims(y, 1)

    total = h.shape[0]
    correct = len(np.where(h==y)[0])
    accuracy = correct / total

    return accuracy

def optimizer(optim_name, epsilon=None, gamma=None):
    if optim_name == 'SGD':
        optim = SGD(gamma=gamma, epsilon=epsilon)
    elif optim_name == 'Momentum':
        optim = Momentum(gamma=gamma, epsilon=epsilon)
    elif optim_name == 'RMSProp':
        optim = RMSProp(gamma=gamma, epsilon=epsilon)
    else:
        raise NotImplementedError
    return optim

config = {
    'Concrete': ('concrete', LinearRegression, RMSE),
    'Graduate': ('graduate', LinearRegression, RMSE),
    'Titanic': ('titanic', LogisticRegression, Accuracy),
    'Digit': ('digit', LogisticRegression, Accuracy)
}

def _initialize(data_name):
    dir_name, model, metric = config[data_name]
    path = os.path.join('./data', dir_name)

    if data_name == 'Concrete':
        train_x, train_y = ConcreteData(path, 'train.csv')
        test_x, test_y = ConcreteData(path, 'test.csv')
    elif data_name == 'Graduate':
        train_x, train_y = GraduateData(path, 'train.csv')
        test_x, test_y = GraduateData(path, 'test.csv')
    elif data_name == 'Digit':
        train_x, train_y = DigitData(path, 'train.csv')
        test_x, test_y = DigitData(path, 'test.csv')
    elif data_name == 'Titanic':
        train_x, train_y = TitanicData(path, 'train.csv')
        test_x, test_y = TitanicData(path, 'test.csv')
    else:
        raise NotImplementedError

    return (train_x, train_y), (test_x, test_y), model, metric

if __name__ == '__main__':
    print('DATASET TEST START\n')
    graduate_x, graduate_y = GraduateData(os.path.join('./data', 'graduate'), 'train.csv')
    concrete_x, concrete_y = ConcreteData(os.path.join('./data', 'concrete'), 'train.csv')
    digit_x, digit_y = DigitData(os.path.join('./data', 'digit'), 'train.csv')
    titanic_x, titanic_y = TitanicData(os.path.join('./data', 'titanic'), 'train.csv')

    print('GRADUATE X, Y : ', graduate_x.shape, ',', graduate_y.shape)
    print('CONCRETE X, Y : ', concrete_x.shape, ',', concrete_y.shape)
    print('DIGIT X, Y : ', digit_x.shape, ',', digit_y.shape)
    print('TITANIC X, Y : ', titanic_x.shape, ',', titanic_y.shape)

    print('\nDATASET TEST FINISHED\n')

    print('OPTIMIZER TEST START\n')

    SGD = optimizer('SGD', 1, 1)
    Momentum = optimizer('Momentum', 1, 1)
    RMSProp = optimizer('RMSProp', 1, 1)

    print('OPTIMIZER TEST FINISHED\n')


import torch
import random
import numpy as np
from GPT2.model import GPT2LMHeadModel
import GPT2.utils as utils
from GPT2.config import GPT2Config
from GPT2.sample import sample_sequence, predict_next
from GPT2.encoder import get_encoder

from typing import Any, Dict, Optional, Generator, List

class TextGenerator:
    def __init__(self, state_dict: Dict[str, Any],
                 seed: Optional[int] = None,
                 ) -> None:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  # type: ignore

        # Load Model
        config = GPT2Config()
        enc = get_encoder()
        model = GPT2LMHeadModel(config)
        model = utils.load_weight(model, state_dict)
        model.to(device)
        model.eval()

        self.device = device
        self.model = model
        self.config = config
        self.enc = enc

        self.start_seed(seed)

    def start_seed(self, seed: Optional[int]) -> None:
        if seed is None:
            seed = random.randint(0, 2147483647)
        np.random.seed(seed)
        torch.random.manual_seed(seed)
        torch.cuda.manual_seed(seed)  # type: ignore

    def generate(
        self,
        text: Optional[str] = None,
        nsamples: int = 1,
        unconditional: bool = False,
        batch_size: Optional[int] = None,
        length: Optional[int] = None,
        temperature: float = 0.7,
        top_k: int = 40,
        quiet: bool = True,
    ) -> Generator[str, None, None]:
        assert (text is None and unconditional) or (text is not None and not unconditional)
        if batch_size is None:
            batch_size = 1
        assert nsamples % batch_size == 0

        if length is not None and length > self.config.n_ctx:
            raise ValueError(f"Can't get samples longer than window size: {self.config.n_ctx}")

        if text is not None:
            context_tokens = self.enc.encode(text)
        else:
            eof = self.enc.encoder['<|endoftext|>']

        for _ in range(nsamples // batch_size):
            out = sample_sequence(
                model=self.model, length=length,
                context=None if unconditional else context_tokens,
                start_token=eof if unconditional else None,
                batch_size=batch_size,
                temperature=temperature, top_k=top_k, device=self.device,
                quiet=quiet
            )
            if unconditional:
                new_seq = out[:, len('<|endoftext|>'):].tolist()
            else:
                new_seq = out[:, len(context_tokens):].tolist()

            for i in range(batch_size):
                yield self.enc.decode(new_seq[i])

    def generate_next_options(
        self,
        text: str,
        temperature: float = 1,
        top_k: int = 0,
        length: Optional[int] = None
    ) -> List[str]:
        context_tokens = self.enc.encode(text)
        out = predict_next(
            self.model, context=context_tokens, temperature=temperature,
            top_k=top_k,
            length=1 if length is None else length,
            device=self.device
        )
        return [self.enc.decode(opt) for opt in out.tolist()]