# -*- coding: utf-8 -*-
"""
.. _tut-overview:

Overview of MEG/EEG analysis with MNE-Python
============================================

This tutorial covers the basic EEG/MEG pipeline for event-related analysis:
loading data, epoching, averaging, plotting, and estimating cortical activity
from sensor data. It introduces the core MNE-Python data structures
:class:`~mne.io.Raw`, :class:`~mne.Epochs`, :class:`~mne.Evoked`, and
:class:`~mne.SourceEstimate`, and covers a lot of ground fairly quickly (at the
expense of depth). Subsequent tutorials address each of these topics in greater
detail.

.. contents:: Page contents
   :local:
   :depth: 1

We begin by importing the necessary Python modules:
"""

import os
import numpy as np
import mne

###############################################################################
# Loading data
# ^^^^^^^^^^^^
#
# MNE-Python data structures are based around the FIF file format from
# Neuromag, but there are reader functions for :ref:`a wide variety of other
# data formats <data-formats>`. MNE-Python also has interfaces to a
# variety of :ref:`publicly available datasets <datasets>`,
# which MNE-Python can download and manage for you.
#
# We'll start this tutorial by loading one of the example datasets (called
# ":ref:`sample-dataset`"), which contains EEG and MEG data from one subject
# performing an audiovisual experiment, along with structural MRI scans for
# that subject. The :func:`mne.datasets.sample.data_path` function will
# automatically download the dataset if it isn't found in one of the expected
# locations, then return the directory path to the dataset (see the
# documentation of :func:`~mne.datasets.sample.data_path` for a list of places
# it checks before downloading). Note also that for this tutorial to run
# smoothly on our servers, we're using a filtered and downsampled version of
# the data (:file:`sample_audvis_filt-0-40_raw.fif`), but an unfiltered version
# (:file:`sample_audvis_raw.fif`) is also included in the sample dataset and
# could be substituted here when running the tutorial locally.

sample_data_folder = mne.datasets.sample.data_path()
sample_data_raw_file = os.path.join(sample_data_folder, 'MEG', 'sample',
                                    'sample_audvis_filt-0-40_raw.fif')
raw = mne.io.read_raw_fif(sample_data_raw_file)

###############################################################################
# By default, :func:`~mne.io.read_raw_fif` displays some information about the
# file it's loading; for example, here it tells us that there are four
# "projection items" in the file along with the recorded data; those are
# :term:`SSP projectors <projector>` calculated to remove environmental noise
# from the MEG signals, plus a projector to mean-reference the EEG channels;
# these are discussed in the tutorial :ref:`tut-projectors-background`.
# In addition to the information displayed during loading,
# you can get a glimpse of the basic details of a :class:`~mne.io.Raw` object
# by printing it; even more is available by printing its ``info`` attribute
# (a :class:`dictionary-like object <mne.Info>` that is preserved across
# :class:`~mne.io.Raw`, :class:`~mne.Epochs`, and :class:`~mne.Evoked`
# objects). The ``info`` data structure keeps track of channel locations,
# applied filters, projectors, etc. Notice especially the ``chs`` entry,
# showing that MNE-Python detects different sensor types and handles each
# appropriately. See :ref:`tut-info-class` for more on the :class:`~mne.Info`
# class.

print(raw)
print(raw.info)

###############################################################################
# :class:`~mne.io.Raw` objects also have several built-in plotting methods;
# here we show the power spectral density (PSD) for each sensor type with
# :meth:`~mne.io.Raw.plot_psd`, as well as a plot of the raw sensor traces with
# :meth:`~mne.io.Raw.plot`. In the PSD plot, we'll only plot frequencies below
# 50 Hz (since our data are low-pass filtered at 40 Hz). In interactive Python
# sessions, :meth:`~mne.io.Raw.plot` is interactive and allows scrolling,
# scaling, bad channel marking, annotation, projector toggling, etc.

raw.plot_psd(fmax=50)
raw.plot(duration=5, n_channels=30)

###############################################################################
# Preprocessing
# ^^^^^^^^^^^^^
#
# MNE-Python supports a variety of preprocessing approaches and techniques
# (maxwell filtering, signal-space projection, independent components analysis,
# filtering, downsampling, etc); see the full list of capabilities in the
# :mod:`mne.preprocessing` and :mod:`mne.filter` submodules. Here we'll clean
# up our data by performing independent components analysis
# (:class:`~mne.preprocessing.ICA`); for brevity we'll skip the steps that
# helped us determined which components best capture the artifacts (see
# :ref:`tut-artifact-ica` for a detailed walk-through of that process).

# set up and fit the ICA
ica = mne.preprocessing.ICA(n_components=20, random_state=97, max_iter=800)
ica.fit(raw)
ica.exclude = [1, 2]  # details on how we picked these are omitted here
ica.plot_properties(raw, picks=ica.exclude)

###############################################################################
# Once we're confident about which component(s) we want to remove, we pass them
# as the ``exclude`` parameter and then apply the ICA to the raw signal. The
# :meth:`~mne.preprocessing.ICA.apply` method requires the raw data to be
# loaded into memory (by default it's only read from disk as-needed), so we'll
# use :meth:`~mne.io.Raw.load_data` first. We'll also make a copy of the
# :class:`~mne.io.Raw` object so we can compare the signal before and after
# artifact removal side-by-side:

orig_raw = raw.copy()
raw.load_data()
ica.apply(raw)

# show some frontal channels to clearly illustrate the artifact removal
chs = ['MEG 0111', 'MEG 0121', 'MEG 0131', 'MEG 0211', 'MEG 0221', 'MEG 0231',
       'MEG 0311', 'MEG 0321', 'MEG 0331', 'MEG 1511', 'MEG 1521', 'MEG 1531',
       'EEG 001', 'EEG 002', 'EEG 003', 'EEG 004', 'EEG 005', 'EEG 006',
       'EEG 007', 'EEG 008']
chan_idxs = [raw.ch_names.index(ch) for ch in chs]
orig_raw.plot(order=chan_idxs, start=12, duration=4)
raw.plot(order=chan_idxs, start=12, duration=4)

###############################################################################
# .. _overview-tut-events-section:
#
# Detecting experimental events
# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
#
# The sample dataset includes several :term:`"STIM" channels <stim channel>`
# that recorded electrical
# signals sent from the stimulus delivery computer (as brief DC shifts /
# squarewave pulses). These pulses (often called "triggers") are used in this
# dataset to mark experimental events: stimulus onset, stimulus type, and
# participant response (button press). The individual STIM channels are
# combined onto a single channel, in such a way that voltage
# levels on that channel can be unambiguously decoded as a particular event
# type. On older Neuromag systems (such as that used to record the sample data)
# this summation channel was called ``STI 014``, so we can pass that channel
# name to the :func:`mne.find_events` function to recover the timing and
# identity of the stimulus events.

events = mne.find_events(raw, stim_channel='STI 014')
print(events[:5])  # show the first 5

###############################################################################
# The resulting events array is an ordinary 3-column :class:`NumPy array
# <numpy.ndarray>`, with sample number in the first column and integer event ID
# in the last column; the middle column is usually ignored. Rather than keeping
# track of integer event IDs, we can provide an *event dictionary* that maps
# the integer IDs to experimental conditions or events. In this dataset, the
# mapping looks like this:
#
# .. _sample-data-event-dict-table:
#
# +----------+----------------------------------------------------------+
# | Event ID | Condition                                                |
# +==========+==========================================================+
# | 1        | auditory stimulus (tone) to the left ear                 |
# +----------+----------------------------------------------------------+
# | 2        | auditory stimulus (tone) to the right ear                |
# +----------+----------------------------------------------------------+
# | 3        | visual stimulus (checkerboard) to the left visual field  |
# +----------+----------------------------------------------------------+
# | 4        | visual stimulus (checkerboard) to the right visual field |
# +----------+----------------------------------------------------------+
# | 5        | smiley face (catch trial)                                |
# +----------+----------------------------------------------------------+
# | 32       | subject button press                                     |
# +----------+----------------------------------------------------------+

event_dict = {'auditory/left': 1, 'auditory/right': 2, 'visual/left': 3,
              'visual/right': 4, 'smiley': 5, 'buttonpress': 32}

###############################################################################
# Event dictionaries like this one are used when extracting epochs from
# continuous data; the ``/`` character in the dictionary keys allows pooling
# across conditions by requesting partial condition descriptors (i.e.,
# requesting ``'auditory'`` will select all epochs with Event IDs 1 and 2;
# requesting ``'left'`` will select all epochs with Event IDs 1 and 3). An
# example of this is shown in the next section. There is also a convenient
# :func:`~mne.viz.plot_events` function for visualizing the distribution of
# events across the duration of the recording (to make sure event detection
# worked as expected). Here we'll also make use of the :class:`~mne.Info`
# attribute to get the sampling frequency of the recording (so our x-axis will
# be in seconds instead of in samples).

fig = mne.viz.plot_events(events, event_id=event_dict, sfreq=raw.info['sfreq'],
                          first_samp=raw.first_samp)

###############################################################################
# For paradigms that are not event-related (e.g., analysis of resting-state
# data), you can extract regularly spaced (possibly overlapping) spans of data
# by creating events using :func:`mne.make_fixed_length_events` and then
# proceeding with epoching as described in the next section.
#
#
# .. _tut-section-overview-epoching:
#
# Epoching continuous data
# ^^^^^^^^^^^^^^^^^^^^^^^^
#
# The :class:`~mne.io.Raw` object and the events array are the bare minimum
# needed to create an :class:`~mne.Epochs` object, which we create with the
# :class:`~mne.Epochs` class constructor. Here we'll also specify some data
# quality constraints: we'll reject any epoch where peak-to-peak signal
# amplitude is beyond reasonable limits for that channel type. This is done
# with a *rejection dictionary*; you may include or omit thresholds for any of
# the channel types present in your data. The values given here are reasonable
# for this particular dataset, but may need to be adapted for different
# hardware or recording conditions. For a more automated approach, consider
# using the `autoreject package`_.

reject_criteria = dict(mag=4000e-15,     # 4000 fT
                       grad=4000e-13,    # 4000 fT/cm
                       eeg=150e-6,       # 150 μV
                       eog=250e-6)       # 250 μV

###############################################################################
# We'll also pass the event dictionary as the ``event_id`` parameter (so we can
# work with easy-to-pool event labels instead of the integer event IDs), and
# specify ``tmin`` and ``tmax`` (the time relative to each event at which to
# start and end each epoch). As mentioned above, by default
# :class:`~mne.io.Raw` and :class:`~mne.Epochs` data aren't loaded into memory
# (they're accessed from disk only when needed), but here we'll force loading
# into memory using the ``preload=True`` parameter so that we can see the
# results of the rejection criteria being applied:

epochs = mne.Epochs(raw, events, event_id=event_dict, tmin=-0.2, tmax=0.5,
                    reject=reject_criteria, preload=True)

###############################################################################
# Next we'll pool across left/right stimulus presentations so we can compare
# auditory versus visual responses. To avoid biasing our signals to the
# left or right, we'll use :meth:`~mne.Epochs.equalize_event_counts` first to
# randomly sample epochs from each condition to match the number of epochs
# present in the condition with the fewest good epochs.

conds_we_care_about = ['auditory/left', 'auditory/right',
                       'visual/left', 'visual/right']
epochs.equalize_event_counts(conds_we_care_about)  # this operates in-place
aud_epochs = epochs['auditory']
vis_epochs = epochs['visual']
del raw, epochs  # free up memory

###############################################################################
# Like :class:`~mne.io.Raw` objects, :class:`~mne.Epochs` objects also have a
# number of built-in plotting methods. One is :meth:`~mne.Epochs.plot_image`,
# which shows each epoch as one row of an image map, with color representing
# signal magnitude; the average evoked response and the sensor location are
# shown below the image:

aud_epochs.plot_image(picks=['MEG 1332', 'EEG 021'])

##############################################################################
# .. note::
#
#     Both :class:`~mne.io.Raw` and :class:`~mne.Epochs` objects have
#     :meth:`~mne.Epochs.get_data` methods that return the underlying data
#     as a :class:`NumPy array <numpy.ndarray>`. Both methods have a ``picks``
#     parameter for subselecting which channel(s) to return; ``raw.get_data()``
#     has additional parameters for restricting the time domain. The resulting
#     matrices have dimension ``(n_channels, n_times)`` for
#     :class:`~mne.io.Raw` and ``(n_epochs, n_channels, n_times)`` for
#     :class:`~mne.Epochs`.
#
# Time-frequency analysis
# ^^^^^^^^^^^^^^^^^^^^^^^
#
# The :mod:`mne.time_frequency` submodule provides implementations of several
# algorithms to compute time-frequency representations, power spectral density,
# and cross-spectral density. Here, for example, we'll compute for the auditory
# epochs the induced power at different frequencies and times, using Morlet
# wavelets. On this dataset the result is not especially informative (it just
# shows the evoked "auditory N100" response); see :ref:`here
# <inter-trial-coherence>` for a more extended example on a dataset with richer
# frequency content.

frequencies = np.arange(7, 30, 3)
power = mne.time_frequency.tfr_morlet(aud_epochs, n_cycles=2, return_itc=False,
                                      freqs=frequencies, decim=3)
power.plot(['MEG 1332'])

###############################################################################
# Estimating evoked responses
# ^^^^^^^^^^^^^^^^^^^^^^^^^^^
#
# Now that we have our conditions in ``aud_epochs`` and ``vis_epochs``, we can
# get an estimate of evoked responses to auditory versus visual stimuli by
# averaging together the epochs in each condition. This is as simple as calling
# the :meth:`~mne.Epochs.average` method on the :class:`~mne.Epochs` object,
# and then using a function from the :mod:`mne.viz` module to compare the
# global field power for each sensor type of the two :class:`~mne.Evoked`
# objects:

aud_evoked = aud_epochs.average()
vis_evoked = vis_epochs.average()

mne.viz.plot_compare_evokeds(dict(auditory=aud_evoked, visual=vis_evoked),
                             legend='upper left', show_sensors='upper right')

###############################################################################
# We can also get a more detailed view of each :class:`~mne.Evoked` object
# using other plotting methods such as :meth:`~mne.Evoked.plot_joint` or
# :meth:`~mne.Evoked.plot_topomap`. Here we'll examine just the EEG channels,
# and see the classic auditory evoked N100-P200 pattern over dorso-frontal
# electrodes, then plot scalp topographies at some additional arbitrary times:

# sphinx_gallery_thumbnail_number = 13
aud_evoked.plot_joint(picks='eeg')
aud_evoked.plot_topomap(times=[0., 0.08, 0.1, 0.12, 0.2], ch_type='eeg')

##############################################################################
# Evoked objects can also be combined to show contrasts between conditions,
# using the :func:`mne.combine_evoked` function. A simple difference can be
# generated by negating one of the :class:`~mne.Evoked` objects passed into the
# function. We'll then plot the difference wave at each sensor using
# :meth:`~mne.Evoked.plot_topo`:

evoked_diff = mne.combine_evoked([aud_evoked, -vis_evoked], weights='equal')
evoked_diff.pick_types('mag').plot_topo(color='r', legend=False)

##############################################################################
# Inverse modeling
# ^^^^^^^^^^^^^^^^
#
# Finally, we can estimate the origins of the evoked activity by projecting the
# sensor data into this subject's :term:`source space` (a set of points either
# on the cortical surface or within the cortical volume of that subject, as
# estimated by structural MRI scans). MNE-Python supports lots of ways of doing
# this (dynamic statistical parametric mapping, dipole fitting, beamformers,
# etc.); here we'll use minimum-norm estimation (MNE) to generate a continuous
# map of activation constrained to the cortical surface. MNE uses a linear
# :term:`inverse operator` to project EEG+MEG sensor measurements into the
# source space. The inverse operator is computed from the
# :term:`forward solution` for this subject and an estimate of :ref:`the
# covariance of sensor measurements <tut_compute_covariance>`. For this
# tutorial we'll skip those computational steps and load a pre-computed inverse
# operator from disk (it's included with the :ref:`sample data
# <sample-dataset>`). Because this "inverse problem" is underdetermined (there
# is no unique solution), here we further constrain the solution by providing a
# regularization parameter specifying the relative smoothness of the current
# estimates in terms of a signal-to-noise ratio (where "noise" here is akin to
# baseline activity level across all of cortex).

# load inverse operator
inverse_operator_file = os.path.join(sample_data_folder, 'MEG', 'sample',
                                     'sample_audvis-meg-oct-6-meg-inv.fif')
inv_operator = mne.minimum_norm.read_inverse_operator(inverse_operator_file)
# set signal-to-noise ratio (SNR) to compute regularization parameter (λ²)
snr = 3.
lambda2 = 1. / snr ** 2
# generate the source time course (STC)
stc = mne.minimum_norm.apply_inverse(vis_evoked, inv_operator,
                                     lambda2=lambda2,
                                     method='MNE')  # or dSPM, sLORETA, eLORETA

##############################################################################
# Finally, in order to plot the source estimate on the subject's cortical
# surface we'll also need the path to the sample subject's structural MRI files
# (the ``subjects_dir``):

# path to subjects' MRI files
subjects_dir = os.path.join(sample_data_folder, 'subjects')
# plot
stc.plot(initial_time=0.1, hemi='split', views=['lat', 'med'],
         subjects_dir=subjects_dir)

##############################################################################
# The remaining tutorials have *much more detail* on each of these topics (as
# well as many other capabilities of MNE-Python not mentioned here:
# connectivity analysis, encoding/decoding models, lots more visualization
# options, etc). Read on to learn more!
#
# .. LINKS
#
# .. _`autoreject package`: http://autoreject.github.io/


# Importing packages and loading in the data set
from utils_pos import get_word_tag, preprocess
import pandas as pd
from collections import defaultdict
import math
import numpy as np

# load in the training corpus
with open("WSJ_02-21.pos", 'r') as f:
    training_corpus = f.readlines()

print(f"A few items of the training corpus list")
print(training_corpus[0:5])

# read the vocabulary data, split by each line of text, and save the list
with open("hmm_vocab.txt", 'r') as f:
    voc_l = f.read().split('\n')

print("A few items of the vocabulary list")
print(voc_l[0:50])
print()
print("A few items at the end of the vocabulary list")
print(voc_l[-50:])

# vocab: dictionary that has the index of the corresponding words
vocab = {}

# Get the index of the corresponding words.
for i, word in enumerate(sorted(voc_l)):
    vocab[word] = i

print("Vocabulary dictionary, key is the word, value is a unique integer")
cnt = 0
for k,v in vocab.items():
    print(f"{k}:{v}")
    cnt += 1
    if cnt > 20:
        break

# load in the test corpus
with open("WSJ_24.pos", 'r') as f:
    y = f.readlines()

print("A sample of the test corpus")
print(y[0:10])

#corpus without tags, preprocessed
_, prep = preprocess(vocab, "test.words")

print('The length of the preprocessed test corpus: ', len(prep))
print('This is a sample of the test_corpus: ')
print(prep[0:10])

# Part 1: Parts-of-speech tagging
# UNQ_C1
def create_dictionaries(training_corpus, vocab):
    """
    Input:
        training_corpus: a corpus where each line has a word followed by its tag.
        vocab: a dictionary where keys are words in vocabulary and value is an index
    Output:
        emission_counts: a dictionary where the keys are (tag, word) and the values are the counts
        transition_counts: a dictionary where the keys are (prev_tag, tag) and the values are the counts
        tag_counts: a dictionary where the keys are the tags and the values are the counts
    """

    # initialize the dictionaries using defaultdict
    emission_counts = defaultdict(int)
    transition_counts = defaultdict(int)
    tag_counts = defaultdict(int)

    # Initialize "prev_tag" (previous tag) with the start state, denoted by '--s--'
    prev_tag = '--s--'

    # use 'i' to track the line number in the corpus
    i = 0

    # Each item in the training corpus contains a word and its POS tag
    # Go through each word and its tag in the training corpus
    for word_tag in training_corpus:

        # Increment the word_tag count
        i += 1

        # Every 50,000 words, print the word count
        if i % 50000 == 0:
            print(f"word count = {i}")

        # get the word and tag using the get_word_tag helper function (imported from utils_pos.py)
        word, tag = get_word_tag(word_tag,vocab)

        # Increment the transition count for the previous word and tag
        transition_counts[(prev_tag, tag)] += 1

        # Increment the emission count for the tag and word
        emission_counts[(tag, word)] += 1

        # Increment the tag count
        tag_counts[tag] += 1

        # Set the previous tag to this tag (for the next iteration of the loop)
        prev_tag = tag

    return emission_counts, transition_counts, tag_counts

  emission_counts, transition_counts, tag_counts = create_dictionaries(training_corpus, vocab)

  # get all the POS states
states = sorted(tag_counts.keys())
print(f"Number of POS tags (number of 'states'): {len(states)}")
print("View these POS tags (states)")
print(states)

print("transition examples: ")
for ex in list(transition_counts.items())[:3]:
    print(ex)
print()

print("emission examples: ")
for ex in list(emission_counts.items())[200:203]:
    print (ex)
print()

print("ambiguous word example: ")
for tup,cnt in emission_counts.items():
    if tup[1] == 'back': print (tup, cnt)

# UNQ_C2 (UNIQUE CELL IDENTIFIER, DO NOT EDIT)
# GRADED FUNCTION: predict_pos

def predict_pos(prep, y, emission_counts, vocab, states):
    '''
    Input:
        prep: a preprocessed version of 'y'. A list with the 'word' component of the tuples.
        y: a corpus composed of a list of tuples where each tuple consists of (word, POS)
        emission_counts: a dictionary where the keys are (tag,word) tuples and the value is the count
        vocab: a dictionary where keys are words in vocabulary and value is an index
        states: a sorted list of all possible tags for this assignment
    Output:
        accuracy: Number of times you classified a word correctly
    '''

    # Initialize the number of correct predictions to zero
    num_correct = 0

    # Get the (tag, word) tuples, stored as a set
    all_words = set(emission_counts.keys())

    # Get the number of (word, POS) tuples in the corpus 'y'
    total = len(y)
    for word, y_tup in zip(prep, y):

        # Split the (word, POS) string into a list of two items
        y_tup_l = y_tup.split()

        # Verify that y_tup contain both word and POS
        if len(y_tup_l) == 2:

            # Set the true POS label for this word
            true_label = y_tup_l[1]

        else:
            # If the y_tup didn't contain word and POS, go to next word
            continue

        count_final = 0
        pos_final = ''

        # If the word is in the vocabulary...
        if word in vocab:
            for pos in states:

            ### START CODE HERE (Replace instances of 'None' with your code) ###

                # define the key as the tuple containing the POS and word
                key = (pos, word)

                # check if the (pos, word) key exists in the emission_counts dictionary
                if key in emission_counts.keys(): # complete this line

                # get the emission count of the (pos,word) tuple
                    count = emission_counts.get(key,0)

                    # keep track of the POS with the largest count
                    if count>count_final: # complete this line

                        # update the final count (largest count)
                        count_final = count

                        # update the final POS
                        pos_final = pos

            # If the final POS (with the largest count) matches the true POS:
            if pos_final==true_label: # complete this line

                # Update the number of correct predictions
                num_correct += 1

    ### END CODE HERE ###
    accuracy = num_correct / total

    return accuracy

accuracy_predict_pos = predict_pos(prep, y, emission_counts, vocab, states)
print(f"Accuracy of prediction using predict_pos is {accuracy_predict_pos:.4f}")

# UNQ_C3
def create_transition_matrix(alpha, tag_counts, transition_counts):
    '''
    Input:
        alpha: number used for smoothing
        tag_counts: a dictionary mapping each tag to its respective count
        transition_counts: transition count for the previous word and tag
    Output:
        A: matrix of dimension (num_tags,num_tags)
    '''
    # Get a sorted list of unique POS tags
    all_tags = sorted(tag_counts.keys())

    # Count the number of unique POS tags
    num_tags = len(all_tags)

    # Initialize the transition matrix 'A'
    A = np.zeros((num_tags,num_tags))

    # Get the unique transition tuples (previous POS, current POS)
    trans_keys = set(transition_counts.keys())

    # Go through each row of the transition matrix A
    for i in range(num_tags):

        # Go through each column of the transition matrix A
        for j in range(num_tags):

            # Initialize the count of the (prev POS, current POS) to zero
            count = 0

            # Define the tuple (prev POS, current POS)
            # Get the tag at position i and tag at position j (from the all_tags list)
            key = (all_tags[i], all_tags[j])

            # Check if the (prev POS, current POS) tuple
            # exists in the transition counts dictionary
            if key in trans_keys: #complete this line

                # Get count from the transition_counts dictionary
                # for the (prev POS, current POS) tuple
                count = transition_counts.get(key,0)

            # Get the count of the previous tag (index position i) from tag_counts
            count_prev_tag = tag_counts.get(all_tags[i],0)

            # Apply smoothing using count of the tuple, alpha,
            # count of previous tag, alpha, and total number of tags
            A[i,j] = (count + alpha)/(count_prev_tag+num_tags*alpha)

    return A

alpha = 0.001
A = create_transition_matrix(alpha, tag_counts, transition_counts)
# Testing your function
print(f"A at row 0, col 0: {A[0,0]:.9f}")
print(f"A at row 3, col 1: {A[3,1]:.4f}")

print("View a subset of transition matrix A")
A_sub = pd.DataFrame(A[30:35,30:35], index=states[30:35], columns = states[30:35] )
print(A_sub)

# UNQ_C4
def create_emission_matrix(alpha, tag_counts, emission_counts, vocab):
    '''
    Input:
        alpha: tuning parameter used in smoothing
        tag_counts: a dictionary mapping each tag to its respective count
        emission_counts: a dictionary where the keys are (tag, word) and the values are the counts
        vocab: a dictionary where keys are words in vocabulary and value is an index.
               within the function it'll be treated as a list
    Output:
        B: a matrix of dimension (num_tags, len(vocab))
    '''

    # get the number of POS tag
    num_tags = len(tag_counts)

    # Get a list of all POS tags
    all_tags = sorted(tag_counts.keys())

    # Get the total number of unique words in the vocabulary
    num_words = len(vocab)

    # Initialize the emission matrix B with places for
    # tags in the rows and words in the columns
    B = np.zeros((num_tags, num_words))

    # Get a set of all (POS, word) tuples
    # from the keys of the emission_counts dictionary
    emis_keys = set(list(emission_counts.keys()))

    # Go through each row (POS tags)
    for i in range(num_tags): # complete this line

        # Go through each column (words)
        for j in range(num_words): # complete this line

            # Initialize the emission count for the (POS tag, word) to zero
            count = 0

            # Define the (POS tag, word) tuple for this row and column
            key =  (all_tags[i], vocab[j])

            # check if the (POS tag, word) tuple exists as a key in emission counts
            if key in emis_keys: # complete this line

                # Get the count of (POS tag, word) from the emission_counts d
                count = emission_counts[key]

            # Get the count of the POS tag
            count_tag = tag_counts[all_tags[i]]

            # Apply smoothing and store the smoothed value
            # into the emission matrix B for this row and column
            B[i,j] = (count+alpha)/(count_tag+num_words*alpha)

    return B

# creating your emission probability matrix. this takes a few minutes to run.
B = create_emission_matrix(alpha, tag_counts, emission_counts, list(vocab))

print(f"View Matrix position at row 0, column 0: {B[0,0]:.9f}")
print(f"View Matrix position at row 3, column 1: {B[3,1]:.9f}")

# Try viewing emissions for a few words in a sample dataframe
cidx  = ['725','adroitly','engineers', 'promoted', 'synergy']

# Get the integer ID for each word
cols = [vocab[a] for a in cidx]

# Choose POS tags to show in a sample dataframe
rvals =['CD','NN','NNS', 'VB','RB','RP']

# For each POS tag, get the row number from the 'states' list
rows = [states.index(a) for a in rvals]

# Get the emissions for the sample of words, and the sample of POS tags
B_sub = pd.DataFrame(B[np.ix_(rows,cols)], index=rvals, columns = cidx )
print(B_sub)

# Part 3: Viterbi Algorithm and Dynamic Programming
# UNQ_C5
def initialize(states, tag_counts, A, B, corpus, vocab):
    '''
    Input:
        states: a list of all possible parts-of-speech
        tag_counts: a dictionary mapping each tag to its respective count
        A: Transition Matrix of dimension (num_tags, num_tags)
        B: Emission Matrix of dimension (num_tags, len(vocab))
        corpus: a sequence of words whose POS is to be identified in a list
        vocab: a dictionary where keys are words in vocabulary and value is an index
    Output:
        best_probs: matrix of dimension (num_tags, len(corpus)) of floats
        best_paths: matrix of dimension (num_tags, len(corpus)) of integers
    '''
    # Get the total number of unique POS tags
    num_tags = len(tag_counts)

    # Initialize best_probs matrix
    # POS tags in the rows, number of words in the corpus as the columns
    best_probs = np.zeros((num_tags, len(corpus)))

    # Initialize best_paths matrix
    # POS tags in the rows, number of words in the corpus as columns
    best_paths = np.zeros((num_tags, len(corpus)), dtype=int)

    # Define the start token
    s_idx = states.index("--s--")

    # Go through each of the POS tags
    for i in range(num_tags): # complete this line

        # Handle the special case when the transition from start token to POS tag i is zero
        if A[s_idx,i]==0: # complete this line

            # Initialize best_probs at POS tag 'i', column 0, to negative infinity
            best_probs[i,0] = float('-inf')

        # For all other cases when transition from start token to POS tag i is non-zero:
        else:

            # Initialize best_probs at POS tag 'i', column 0
            # Check the formula in the instructions above
            best_probs[i,0] = math.log(A[s_idx,i]) + math.log(B[i,vocab[corpus[0]]])

    return best_probs, best_paths

best_probs, best_paths = initialize(states, tag_counts, A, B, prep, vocab)
# Test the function
print(f"best_probs[0,0]: {best_probs[0,0]:.4f}")
print(f"best_paths[2,3]: {best_paths[2,3]:.4f}")

# UNQ_C6
def viterbi_forward(A, B, test_corpus, best_probs, best_paths, vocab):
    '''
    Input:
        A, B: The transition and emission matrices respectively
        test_corpus: a list containing a preprocessed corpus
        best_probs: an initilized matrix of dimension (num_tags, len(corpus))
        best_paths: an initilized matrix of dimension (num_tags, len(corpus))
        vocab: a dictionary where keys are words in vocabulary and value is an index
    Output:
        best_probs: a completed matrix of dimension (num_tags, len(corpus))
        best_paths: a completed matrix of dimension (num_tags, len(corpus))
    '''
    # Get the number of unique POS tags (which is the num of rows in best_probs)
    num_tags = best_probs.shape[0]

    # Go through every word in the corpus starting from word 1
    # Recall that word 0 was initialized in `initialize()`
    for i in range(1, len(test_corpus)):

        # Print number of words processed, every 5000 words
        if i % 5000 == 0:
            print("Words processed: {:>8}".format(i))

        # For each unique POS tag that the current word can be
        for j in range(num_tags): # complete this line

            # Initialize best_prob for word i to negative infinity
            best_prob_i = float('-inf')

            # Initialize best_path for current word i to None
            best_path_i = None

            # For each POS tag that the previous word can be:
            for k in range(num_tags): # complete this line

                # Calculate the probability =
                # best probs of POS tag k, previous word i-1 +
                # log(prob of transition from POS k to POS j) +
                # log(prob that emission of POS j is word i)
                prob = best_probs[k,i-1] + math.log(A[k,j]) + math.log(B[j,vocab[test_corpus[i]]])

                # check if this path's probability is greater than
                # the best probability up to and before this point
                if prob > best_prob_i: # complete this line

                    # Keep track of the best probability
                    best_prob_i = prob

                    # keep track of the POS tag of the previous word
                    # that is part of the best path.
                    # Save the index (integer) associated with
                    # that previous word's POS tag
                    best_path_i = k

            # Save the best probability for the
            # given current word's POS tag
            # and the position of the current word inside the corpus
            best_probs[j,i] = best_prob_i

            # Save the unique integer ID of the previous POS tag
            # into best_paths matrix, for the POS tag of the current word
            # and the position of the current word inside the corpus.
            best_paths[j,i] = best_path_i

    return best_probs, best_paths

# this will take a few minutes to run => processes ~ 30,000 words
best_probs, best_paths = viterbi_forward(A, B, prep, best_probs, best_paths, vocab)

# Test this function
print(f"best_probs[0,1]: {best_probs[0,1]:.4f}")
print(f"best_probs[0,4]: {best_probs[0,4]:.4f}")

# UNQ_C7
# UNQ_C7 (UNIQUE CELL IDENTIFIER, DO NOT EDIT)
# GRADED FUNCTION: viterbi_backward
def viterbi_backward(best_probs, best_paths, corpus, states):
    '''
    This function returns the best path.

    '''
    # Get the number of words in the corpus
    # which is also the number of columns in best_probs, best_paths
    m = best_paths.shape[1]

    # Initialize array z, same length as the corpus
    z = [None] * m

    # Get the number of unique POS tags
    num_tags = best_probs.shape[0]

    # Initialize the best probability for the last word
    best_prob_for_last_word = float('-inf')

    # Initialize pred array, same length as corpus
    pred = [None] * m

    ## Step 1 ##

    # Go through each POS tag for the last word (last column of best_probs)
    # in order to find the row (POS tag integer ID)
    # with highest probability for the last word
    for k in range(num_tags): # complete this line

        # If the probability of POS tag at row k
        # is better than the previously best probability for the last word:
        if best_probs[k,-1] > best_prob_for_last_word: # complete this line

            # Store the new best probability for the last word
            best_prob_for_last_word = best_probs[k,m-1]

            # Store the unique integer ID of the POS tag
            # which is also the row number in best_probs
            z[m - 1] = k

    # Convert the last word's predicted POS tag
    # from its unique integer ID into the string representation
    # using the 'states' list
    # store this in the 'pred' array for the last word
    pred[m - 1] = states[k]

    ## Step 2 ##
    # Find the best POS tags by walking backward through the best_paths
    # From the last word in the corpus to the 0th word in the corpus
    for i in range(m-1, -1, -1): # complete this line

        # Retrieve the unique integer ID of
        # the POS tag for the word at position 'i' in the corpus
        pos_tag_for_word_i = z[i]

        # In best_paths, go to the row representing the POS tag of word i
        # and the column representing the word's position in the corpus
        # to retrieve the predicted POS for the word at position i-1 in the corpus
        z[i - 1] = best_paths[pos_tag_for_word_i,i]

        # Get the previous word's POS tag in string form
        # Use the 'states' list,
        # where the key is the unique integer ID of the POS tag,
        # and the value is the string representation of that POS tag
        pred[i - 1] = states[z[i - 1]]

    return pred

# Run and test your function
pred = viterbi_backward(best_probs, best_paths, prep, states)
m=len(pred)
print('The prediction for pred[-7:m-1] is: \n', prep[-7:m-1], "\n", pred[-7:m-1], "\n")
print('The prediction for pred[0:8] is: \n', pred[0:7], "\n", prep[0:7])

print('The third word is:', prep[3])
print('Your prediction is:', pred[3])
print('Your corresponding label y is: ', y[3])

# UNQ_C8
# UNQ_C8 (UNIQUE CELL IDENTIFIER, DO NOT EDIT)
# GRADED FUNCTION: compute_accuracy
def compute_accuracy(pred, y):
    '''
    Input:
        pred: a list of the predicted parts-of-speech
        y: a list of lines where each word is separated by a '\t' (i.e. word \t tag)
    Output:

    '''
    num_correct = 0
    total = 0

    # Zip together the prediction and the labels
    for prediction, y in zip(pred, y):
        # Split the label into the word and the POS tag
        word_tag_tuple = y.split()

        # Check that there is actually a word and a tag
        # no more and no less than 2 items
        if len(word_tag_tuple) != 2: # complete this line
            continue

        #print(word_tag_tuple)
        # store the word and tag separately
        word, tag = word_tag_tuple

        # Check if the POS tag label matches the prediction
        if prediction==tag: # complete this line

            # count the number of times that the prediction
            # and label match
            num_correct += 1

        # keep track of the total number of examples (that have valid labels)
        total += 1

    return num_correct/total

print(f"Accuracy of the Viterbi algorithm is {compute_accuracy(pred, y):.4f}")


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

__author__ = 'qiudebo'

import numpy as np
from matplotlib import pyplot as plt

if __name__ == '__main__':
    plt.figure()
    plt.grid(True)

    data = np.linspace(1, 10, 50)
    plt.plot(data)

    plt.show()


"""Verify that simulation produced by ANUGA compares to published
validation timeseries ch5, ch7 and ch9 as well as the boundary timeseries.

RMS norm is printed and plots are produced as png files.
No plots are shown on screen.
"""

import sys

import numpy as num
import anuga
from anuga.file.netcdf import NetCDFFile

import project
from anuga.abstract_2d_finite_volumes.util import file_function
from anuga.utilities.numerical_tools import ensure_numeric
from anuga.utilities.numerical_tools import cov
#from anuga.utilities.numerical_tools import get_machine_precision
from anuga.shallow_water.sww_interrogate import get_maximum_inundation_elevation
from anuga.shallow_water.sww_interrogate import get_maximum_inundation_location

args = anuga.get_args()
verbose = args.verbose

testing = True

#-------------------------
# Basic data
#-------------------------

finaltime = 22.5
timestep = 0.05

gauge_locations = [[0.000, 1.696]] # Boundary gauge
gauge_locations += [[4.521, 1.196],  [4.521, 1.696],  [4.521, 2.196]] #Ch 5-7-9
gauge_names = ['Boundary', 'ch5', 'ch7', 'ch9']

validation_data = {}
for key in gauge_names:
    validation_data[key] = []

# Expected values
expected_covariance = {'Boundary': 5.269569575007607815e-05,
                       'ch5': 1.166277999581819919e-04,
                       'ch7': 1.127136457890861503e-04,
                       'ch9': 1.250659477418482129e-04}

expected_difference = {'Boundary': 8.350712673810733924e-04,
                       'ch5': 3.405426180525532483e-03,
                       'ch7': 2.852870417368218517e-03,
                       'ch9': 3.248778982037564891e-03}

expected_maximum = {'Boundary': 1.611749508386188523e-02,
                    'ch5': 3.551308418158714147e-02,
                    'ch7': 3.858418457126511908e-02,
                    'ch9': 4.317962986578308127e-02}

expected_minimum = {'Boundary': -1.164547474575844919e-02,
                    'ch5': -8.664439185502026408e-03,
                    'ch7': -2.726335488279797541e-03,
                    'ch9': -5.977581218447349659e-03}

expected_argmax = {'Boundary': 1.255000000000000071e+01,
                   'ch5': 1.839999999999999858e+01,
                   'ch7': 1.700000000000000000e+01,
                   'ch9': 1.685000000000000142e+01}

expected_argmin = {'Boundary': 2.064999999999999858e+01,
                   'ch5': 1.459999999999999964e+01,
                   'ch7': 1.230000000000000071e+01,
                   'ch9': 1.315000000000000036e+01}

#-------------------------
# Read validation data
#-------------------------

if verbose: print('Reading', project.boundary_filename)

fid = NetCDFFile(project.boundary_filename, 'r')
input_time = fid.variables['time'][:]
validation_data['Boundary'] = fid.variables['stage'][:]

reference_time = []
fid = open(project.validation_filename)
lines = fid.readlines()
fid.close()

for i, line in enumerate(lines[1:]):
    if i == len(input_time): break

    fields = line.split()

    reference_time.append(float(fields[0]))    # Record reference time
    for j, key in enumerate(gauge_names[1:]):  # Omit boundary gauge
        value = float(fields[1:][j])           # Omit time
        validation_data[key].append(value/100) # Convert cm2m

# Checks
assert reference_time[0] == 0.0
assert reference_time[-1] == finaltime
assert num.allclose(reference_time, input_time)

for key in gauge_names:
    validation_data[key] = ensure_numeric(validation_data[key])

#--------------------------------------------------
# Read and interpolate model output
#--------------------------------------------------

#if len(sys.argv) > 1:
#    sww_filename = sys.argv[1]
#else:
sww_filename = project.output_filename

f = file_function(sww_filename,
                  quantities='stage',
                  interpolation_points=gauge_locations,
                  use_cache=False,
                  verbose=verbose)

def report_difference(name, computed_value, reference_value, rtol, atol):

    if abs(reference_value) > 0:
        msg = '%s (expected, computed):\n  (%.18e, %.18e):\n  Relative error=%.18e'\
              %(name, reference_value, computed_value,
                abs(reference_value-computed_value)/reference_value)
        print(msg)

    msg = '  Absolute error=%.18e'\
          %(abs(reference_value-computed_value))
    print(msg)

    #print 'Allclose:', allclose(reference_value, computed_value,
    #                            rtol=rtol, atol=atol)
    if testing is True:
        assert num.allclose(reference_value, computed_value,
                            rtol=rtol, atol=atol), msg

#--------------------------------------------------
# Compare model output to validation data
#--------------------------------------------------

#eps = get_machine_precision()

# Tolerances  for 20,000 triangles
rtol = 2.0e-2
atol = 2.0e-2

# Tolerances  for 60,000 triangles
#rtol = 1.0e-2
#atol = 1.0e-2

if verbose: print('Precisions used: rtol=%e, atol=%e' %(rtol, atol))

#print reference_time
for k, name in enumerate(gauge_names):

    sqsum = 0
    denom = 0
    model = []
    if verbose:
        print()
        print('Validating ' + name)
    observed_timeseries = validation_data[name]
    for i, t in enumerate(reference_time):
        model.append(f(t, point_id=k)[0])

    # Covariance measure
    res = cov(observed_timeseries, model)
    if verbose:
        report_difference('Covariance', res, expected_covariance[name], rtol, atol)

    # Difference measures
    res = sum(abs(observed_timeseries-model))/len(model)
    if verbose:
        report_difference('Accumulated difference', res,
                      expected_difference[name], rtol, atol)

    # Extrema
    res = max(model)
    if verbose:
        report_difference('Maximum', res, expected_maximum[name], rtol, atol)

    res = min(model)
    if verbose:
        report_difference('Minimum', res, expected_minimum[name], rtol, atol)

    # Locations of extrema
    #i0 = argmax(observed_timeseries)
    i1 = num.argmax(model)
    res = reference_time[i1]
    if verbose:
        report_difference('Location of maximum', res, expected_argmax[name], rtol, atol)

    if not name in ['ch7', 'ch9']:
        # Minima of ch7 and ch9 are very flat and hard to pinpoint
        i1 = num.argmin(model)
        res = reference_time[i1]
        if verbose:
            report_difference('Location of minimum', res, expected_argmin[name],
                          rtol, atol)

# Check max runup

q = get_maximum_inundation_elevation(sww_filename)
loc = get_maximum_inundation_location(sww_filename)

if verbose:
    print('Max runup elevation: ', q)
    print('Max runup elevation (scaled by 400): ', q*400)
    print('Max runup location:  ', loc)


def test():
    import inspect
    import numpy as np
    # Here we can either check objects created in the solution code, or the
    # string value of the solution, available as __solution__. A helper for
    # printing formatted messages is available as __msg__. See the testTemplate
    # in the meta.json for details.

    # If an assertion fails, the message will be displayed
    assert "arr2" in __solution__, "Your object does not exist. Please make sure you are naming your object 'arr2'."
    assert "arr2t" in __solution__, "Your object does not exist. Please make sure you are naming your object 'arr2t'."
    assert "sliced_arr2t" in __solution__, "Your object does not exist. Please make sure you are naming your object 'sliced_arr2t'."
    assert arr2.shape == (2, 3), "The dimensions of your array are incorrect. Make sure you are creating a 2D array using the 'reshape()' function."
    assert arr2t.shape == (3, 2), "The dimensions of the transposed array are incorrect. Make sure you are transposing the array properly."
    assert sliced_arr2t.shape == (3,), "The dimensions of the sliced array are incorrect. Make sure you are only slicing the required values."
    assert sum(sliced_arr2t) == 39.0, "The the values in the sliced array are incorrect. Are you slicing properly?"
    __msg__.good("Nice work, well done!")


# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import string

import numpy as np
import pytest

from nemo.collections.common.tokenizers.column_coder import ColumnCodes
from nemo.collections.common.tokenizers.tabular_tokenizer import TabularTokenizer

class TestTabularTokenizer:
    def setup_method(self, test_method):
        column_configs = [
            {
                "name": "col_a",
                "code_type": "float",
                "args": {"code_len": 4, "base": 16, "fillall": False, "hasnan": True, "transform": 'yeo-johnson'},
            },
            {
                "name": "col_b",
                "code_type": "float",
                "args": {"code_len": 4, "base": 177, "fillall": True, "hasnan": True, "transform": 'quantile'},
            },
            {
                "name": "col_c",
                "code_type": "int",
                "args": {"code_len": 3, "base": 12, "fillall": True, "hasnan": True},
            },
            {"name": "col_d", "code_type": "category",},
        ]

        example_arrays = {}
        np.random.seed(1234)

        array = np.random.random(100)
        example_arrays['col_a'] = array

        array = np.random.random(100)
        example_arrays['col_b'] = array

        array = np.random.randint(3, 1000, 100)
        example_arrays['col_c'] = array

        ALPHABET = np.array(list(string.ascii_lowercase + ' '))
        array = np.char.add(np.random.choice(ALPHABET, 1000), np.random.choice(ALPHABET, 1000))
        example_arrays['col_d'] = array

        self.cc = ColumnCodes.get_column_codes(column_configs, example_arrays)

    @pytest.mark.unit
    def test_tabular_tokenizer(self):
        tab = TabularTokenizer(self.cc, delimiter=',')
        text = "0.323, 0.1, 232, xy\n0.323, 0.1, 232, xy<|endoftext|>"
        r = tab.text_to_tokens(text)
        assert len(r) == 10
        assert tab.eod == 1351
        assert tab.eor == 1352
        assert tab.num_columns == 4
        assert self.cc.vocab_size == 1351
        assert tab.vocab_size == 1353
        r = tab.text_to_ids(text)
        assert (sum(self.cc.sizes) + 1) * 2 == len(r)
        assert np.array_equal(
            np.array(r[0:13]), np.array([49, 32, 29, 15, 584, 417, 305, 76, 787, 780, 773, 1313, 1352])
        )
        assert np.array_equal(
            np.array(r[13:]), np.array([49, 32, 29, 15, 584, 417, 305, 76, 787, 780, 773, 1313, 1351])
        )
        reversed_text = tab.ids_to_text(r)
        assert reversed_text == '0.3230,0.0999998,232,xy\n0.3230,0.0999998,232,xy<|endoftext|>'

        text = "xy\n0.323, 0.1, 232, xy<|endoftext|>"
        r = tab.text_to_tokens(text)
        assert len(r) == 7
        r = tab.text_to_ids(text)
        assert sum(self.cc.sizes) + 1 + 2 == len(r)
        assert np.array_equal(np.array(r[0:2]), np.array([1313, 1352]))
        assert np.array_equal(
            np.array(r[2:15]), np.array([49, 32, 29, 15, 584, 417, 305, 76, 787, 780, 773, 1313, 1351])
        )
        reversed_text = tab.ids_to_text(r)
        assert reversed_text == 'xy\n0.3230,0.0999998,232,xy<|endoftext|>'

        text = "\n0.323, 0.1, 232, xy<|endoftext|>"
        r = tab.text_to_tokens(text)
        assert len(r) == 5
        r = tab.text_to_ids(text)
        assert sum(self.cc.sizes) + 1 == len(r)
        assert np.array_equal(
            np.array(r[0:13]), np.array([49, 32, 29, 15, 584, 417, 305, 76, 787, 780, 773, 1313, 1351])
        )
        reversed_text = tab.ids_to_text(r)
        assert reversed_text == '0.3230,0.0999998,232,xy<|endoftext|>'

        text = "232, xy\n0.323, 0.1, 232, xy<|endoftext|>"
        r = tab.text_to_tokens(text)
        assert len(r) == 8
        r = tab.text_to_ids(text)
        assert sum(self.cc.sizes) + 1 + 5 == len(r)
        assert np.array_equal(np.array(r[0:5]), np.array([787, 780, 773, 1313, 1352]))
        assert np.array_equal(
            np.array(r[5:18]), np.array([49, 32, 29, 15, 584, 417, 305, 76, 787, 780, 773, 1313, 1351])
        )
        reversed_text = tab.ids_to_text(r)
        assert reversed_text == '232,xy\n0.3230,0.0999998,232,xy<|endoftext|>'

        text = "0.1, 232, xy\n0.323, 0.1, 232, xy<|endoftext|>"
        r = tab.text_to_tokens(text)
        assert len(r) == 9
        r = tab.text_to_ids(text)
        assert sum(self.cc.sizes) + 1 + 9 == len(r)
        assert np.array_equal(np.array(r[0:9]), np.array([584, 417, 305, 76, 787, 780, 773, 1313, 1352]))
        assert np.array_equal(
            np.array(r[9:22]), np.array([49, 32, 29, 15, 584, 417, 305, 76, 787, 780, 773, 1313, 1351])
        )
        reversed_text = tab.ids_to_text(r)
        assert reversed_text == '0.0999998,232,xy\n0.3230,0.0999998,232,xy<|endoftext|>'


import ec602lib
import unittest
import random
import subprocess
import os
import numpy as np

progname = "polyops.cpp"

valid_includes = set(['vector'])

TIMEALLOWED = 1
COMPILEALLOWED = 2

P={"stdout":subprocess.PIPE,"timeout":TIMEALLOWED,"stderr":subprocess.PIPE}

CP={"stdout":subprocess.PIPE,"timeout":COMPILEALLOWED,"stderr":subprocess.PIPE}

refcode={'lines':29,'words':130}

cppmain="""
#include <iostream>
#include <vector>

using namespace std;

#include "PROGNAME"

int main()
{

  int Alen,Blen;

  cin >> Alen >> Blen;

  Poly A(Alen,0),B(Blen,0);

  for (auto& e : A)
     cin >> e;

  for (auto& e : B)
     cin >> e;

  for (auto e : add_poly(A,B))
     cout << e << " ";
  cout << endl;

  for (auto e : multiply_poly(A,B))
     cout << e << " ";
  cout << endl;

}"""

Same_Add_Tests=[
[(1.2,0,5),(3,2,1),(4.2,2,6)],
[(1,2,3),(4,5,6),(5,7,9)],
[(0,0,4),(0,0,0.1),(0,0,4.1)],
]

Diff_Add_Tests=[
[(1.2,0,0, 5),(3,2,1),(4.2,2,1,5)],
[(1,6,3),(4,),(5,6,3)],
[(0,0,0,0,5),(0,0,1,4),(0,0,1,4,5)],
]

Same_Mul_Tests=[
[(1.2,0,5),(3,2,1),(3.6, 2.4, 16.2, 10, 5)],
[(1,2,3),(4,5,6),(4, 13, 28, 27, 18)],
[(0,0,4),(0,0,0.1),(0, 0, 0, 0, 0.4)],
]

Diff_Mul_Tests=[
[(1.2,0,0, 5),(3,2,1),(3.6, 2.4, 1.2, 15, 10, 5)],
[(1,6,3),(4,),(4, 24, 12)],
[(1,1,-1),(1,1,1),(1, 2, 1, 0.0, -1)],
[(0,0,0,0,5),(0,0,1,4),(0, 0, 0, 0, 0, 0, 5, 20)],
]

Tricky_Add_Tests=[
[(1,3,2),(1,6,-2),(2,9)],
[(1,1,1),(-1,-1,-1),(0,) ],
]

Tricky_Mul_Tests=[
[(1,2,1,1,1,6),(0,),(0,) ],
]

double_add_tests=[
[(1e-50,),(-0.99e-50,),(1e-52,)],
[(1e300,2e300),(0,0,3e300),(1e300,2e300,3e300)],
]

def fin(a,b):
  "format the input for polyops_example_main"
  astr = " ".join(str(x) for x in a)
  bstr = " ".join(str(x) for x in b)
  return "{} {} {} {}".format(len(a),len(b),astr,bstr).encode()

def fout(T):
  "extract the output for polyops_example_main"
  text = T.stdout.decode().splitlines()
  addition = tuple(float(x) for x in text[0].strip().split())
  multiply = tuple(float(x) for x in text[1].strip().split())
  return addition,multiply

def compile(cpp,executable):
  return ['g++','-std=c++14',cpp, '-o', executable]

def check_add(self,a,b,aplusb):
   with self.subTest(CASE=" {} + {} = {}".format(a,b,aplusb)):
      T = subprocess.run([self.executable],input=fin(a,b),**P)
      add_res,_ = fout(T)
      #self.assertEqual(len(add_res),len(aplusb),
      #     "your addition: {}\ncorrect answer: {}\n".format(add_res,aplusb))
      if len(add_res) != len(aplusb) or not np.allclose(add_res,aplusb,atol=0):
        self.fail("your addition: {}\ncorrect answer: {}\n".format(add_res,aplusb))

def check_mult(self,a,b,atimesb):
    with self.subTest(CASE=" {} * {} = {}".format(a,b,atimesb)):
      T = subprocess.run([self.executable],input=fin(a,b),**P)
      _,res = fout(T)
      if len(res) != len(atimesb) or not np.allclose(res,atimesb,atol=0):
        self.fail("your multiply: {}\ncorrect answer: {}\n".format(res,atimesb))

class polyopsTestCase(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        baseprogname = 'polyopsmain'+str(random.randint(1000,100000))
        mainprogname = baseprogname+'.cpp'
        with open(mainprogname,'w') as f:
            f.write(cppmain.replace('PROGNAME',progname))

        try:
          C = subprocess.run(compile(mainprogname,baseprogname),**CP)
        except Exception as e:
          raise unittest.SkipTest("Compile failed.\n"+str(e))
        finally:
          os.remove(mainprogname)

        if C.returncode:
            raise unittest.SkipTest("Compile failed.\n"+str(C.stderr.decode()))

        cls.executable = baseprogname

    @classmethod
    def tearDownClass(cls):
       try:
        os.remove(cls.executable)
       except:
        pass

    def test_includes(self):
        "a. check the included libraries are allowed"
        f=open(progname)
        file_contents=f.read()
        f.close()
        includes = ec602lib.get_includes(file_contents)
        invalid_includes = includes - valid_includes
        if invalid_includes:
          self.fail('Invalid includes: {}'.format(" ".join(x for x in invalid_includes)))

    def test_add_same_size(self):
       "b. add same size vectors"
       for (a,b,res) in Same_Add_Tests:
           check_add(self,a,b,res)

    def test_add_different_size(self):
       "c. add different size vectors"
       for (a,b,res) in Diff_Add_Tests:
           check_add(self,a,b,res)

    def test_mult_same_size(self):
       "d. multiply same size vectors"
       for (a,b,res) in Same_Mul_Tests:
           check_mult(self,a,b,res)

    def test_mult_different_size(self):
       "e. multiply different size vectors"
       for (a,b,res) in Diff_Mul_Tests:
           check_mult(self,a,b,res)

    def test_add_tricky(self):
       "f. add vectors with result smaller"
       for (a,b,res) in Tricky_Add_Tests:
           check_add(self,a,b,res)

    def test_mult_tricky(self):
       "g. multiply vectors with result smaller"
       for (a,b,res) in Tricky_Mul_Tests:
           check_mult(self,a,b,res)
    def test_double_add(self):
       "h. vector double"
       for (a,b,res) in double_add_tests:
             check_add(self,a,b,res)

if __name__=="__main__":
    _,results,_ = ec602lib.overallcpp(progname,polyopsTestCase,refcode,docompile=False,)
    #unittest.main()
    print(results)


"""
Tests specific to `np.loadtxt` added during the move of loadtxt to be backed
by C code.
These tests complement those found in `test_io.py`.
"""

import sys
import pytest
from tempfile import NamedTemporaryFile, mkstemp
from io import StringIO

import numpy as np
from numpy.ma.testutils import assert_equal
from numpy.testing import assert_array_equal, HAS_REFCOUNT, IS_PYPY

def test_scientific_notation():
    """Test that both 'e' and 'E' are parsed correctly."""
    data = StringIO(
        (
            "1.0e-1,2.0E1,3.0\n"
            "4.0e-2,5.0E-1,6.0\n"
            "7.0e-3,8.0E1,9.0\n"
            "0.0e-4,1.0E-1,2.0"
        )
    )
    expected = np.array(
        [[0.1, 20., 3.0], [0.04, 0.5, 6], [0.007, 80., 9], [0, 0.1, 2]]
    )
    assert_array_equal(np.loadtxt(data, delimiter=","), expected)

@pytest.mark.parametrize("comment", ["..", "//", "@-", "this is a comment:"])
def test_comment_multiple_chars(comment):
    content = "# IGNORE\n1.5, 2.5# ABC\n3.0,4.0# XXX\n5.5,6.0\n"
    txt = StringIO(content.replace("#", comment))
    a = np.loadtxt(txt, delimiter=",", comments=comment)
    assert_equal(a, [[1.5, 2.5], [3.0, 4.0], [5.5, 6.0]])

@pytest.fixture
def mixed_types_structured():
    """
    Fixture providing hetergeneous input data with a structured dtype, along
    with the associated structured array.
    """
    data = StringIO(
        (
            "1000;2.4;alpha;-34\n"
            "2000;3.1;beta;29\n"
            "3500;9.9;gamma;120\n"
            "4090;8.1;delta;0\n"
            "5001;4.4;epsilon;-99\n"
            "6543;7.8;omega;-1\n"
        )
    )
    dtype = np.dtype(
        [('f0', np.uint16), ('f1', np.float64), ('f2', 'S7'), ('f3', np.int8)]
    )
    expected = np.array(
        [
            (1000, 2.4, "alpha", -34),
            (2000, 3.1, "beta", 29),
            (3500, 9.9, "gamma", 120),
            (4090, 8.1, "delta", 0),
            (5001, 4.4, "epsilon", -99),
            (6543, 7.8, "omega", -1)
        ],
        dtype=dtype
    )
    return data, dtype, expected

@pytest.mark.parametrize('skiprows', [0, 1, 2, 3])
def test_structured_dtype_and_skiprows_no_empty_lines(
        skiprows, mixed_types_structured):
    data, dtype, expected = mixed_types_structured
    a = np.loadtxt(data, dtype=dtype, delimiter=";", skiprows=skiprows)
    assert_array_equal(a, expected[skiprows:])

def test_unpack_structured(mixed_types_structured):
    data, dtype, expected = mixed_types_structured

    a, b, c, d = np.loadtxt(data, dtype=dtype, delimiter=";", unpack=True)
    assert_array_equal(a, expected["f0"])
    assert_array_equal(b, expected["f1"])
    assert_array_equal(c, expected["f2"])
    assert_array_equal(d, expected["f3"])

def test_structured_dtype_with_shape():
    dtype = np.dtype([("a", "u1", 2), ("b", "u1", 2)])
    data = StringIO("0,1,2,3\n6,7,8,9\n")
    expected = np.array([((0, 1), (2, 3)), ((6, 7), (8, 9))], dtype=dtype)
    assert_array_equal(np.loadtxt(data, delimiter=",", dtype=dtype), expected)

def test_structured_dtype_with_multi_shape():
    dtype = np.dtype([("a", "u1", (2, 2))])
    data = StringIO("0 1 2 3\n")
    expected = np.array([(((0, 1), (2, 3)),)], dtype=dtype)
    assert_array_equal(np.loadtxt(data, dtype=dtype), expected)

def test_nested_structured_subarray():
    # Test from gh-16678
    point = np.dtype([('x', float), ('y', float)])
    dt = np.dtype([('code', int), ('points', point, (2,))])
    data = StringIO("100,1,2,3,4\n200,5,6,7,8\n")
    expected = np.array(
        [
            (100, [(1., 2.), (3., 4.)]),
            (200, [(5., 6.), (7., 8.)]),
        ],
        dtype=dt
    )
    assert_array_equal(np.loadtxt(data, dtype=dt, delimiter=","), expected)

def test_structured_dtype_offsets():
    # An aligned structured dtype will have additional padding
    dt = np.dtype("i1, i4, i1, i4, i1, i4", align=True)
    data = StringIO("1,2,3,4,5,6\n7,8,9,10,11,12\n")
    expected = np.array([(1, 2, 3, 4, 5, 6), (7, 8, 9, 10, 11, 12)], dtype=dt)
    assert_array_equal(np.loadtxt(data, delimiter=",", dtype=dt), expected)

@pytest.mark.parametrize("param", ("skiprows", "max_rows"))
def test_exception_negative_row_limits(param):
    """skiprows and max_rows should raise for negative parameters."""
    with pytest.raises(ValueError, match="argument must be nonnegative"):
        np.loadtxt("foo.bar", **{param: -3})

@pytest.mark.parametrize("param", ("skiprows", "max_rows"))
def test_exception_noninteger_row_limits(param):
    with pytest.raises(TypeError, match="argument must be an integer"):
        np.loadtxt("foo.bar", **{param: 1.0})

@pytest.mark.parametrize(
    "data, shape",
    [
        ("1 2 3 4 5\n", (1, 5)),  # Single row
        ("1\n2\n3\n4\n5\n", (5, 1)),  # Single column
    ]
)
def test_ndmin_single_row_or_col(data, shape):
    arr = np.array([1, 2, 3, 4, 5])
    arr2d = arr.reshape(shape)

    assert_array_equal(np.loadtxt(StringIO(data), dtype=int), arr)
    assert_array_equal(np.loadtxt(StringIO(data), dtype=int, ndmin=0), arr)
    assert_array_equal(np.loadtxt(StringIO(data), dtype=int, ndmin=1), arr)
    assert_array_equal(np.loadtxt(StringIO(data), dtype=int, ndmin=2), arr2d)

@pytest.mark.parametrize("badval", [-1, 3, None, "plate of shrimp"])
def test_bad_ndmin(badval):
    with pytest.raises(ValueError, match="Illegal value of ndmin keyword"):
        np.loadtxt("foo.bar", ndmin=badval)

@pytest.mark.parametrize(
    "ws",
    (
            "\t",  # tab
            "\u2003",  # em
            "\u00A0",  # non-break
            "\u3000",  # ideographic space
    )
)
def test_blank_lines_spaces_delimit(ws):
    txt = StringIO(
        f"1 2{ws}30\n\n4 5 60\n  {ws}  \n7 8 {ws} 90\n  # comment\n3 2 1"
    )
    # NOTE: It is unclear that the `  # comment` should succeed. Except
    #       for delimiter=None, which should use any whitespace (and maybe
    #       should just be implemented closer to Python
    expected = np.array([[1, 2, 30], [4, 5, 60], [7, 8, 90], [3, 2, 1]])
    assert_equal(
        np.loadtxt(txt, dtype=int, delimiter=None, comments="#"), expected
    )

def test_blank_lines_normal_delimiter():
    txt = StringIO('1,2,30\n\n4,5,60\n\n7,8,90\n# comment\n3,2,1')
    expected = np.array([[1, 2, 30], [4, 5, 60], [7, 8, 90], [3, 2, 1]])
    assert_equal(
        np.loadtxt(txt, dtype=int, delimiter=',', comments="#"), expected
    )

@pytest.mark.parametrize("dtype", (float, object))
def test_maxrows_no_blank_lines(dtype):
    txt = StringIO("1.5,2.5\n3.0,4.0\n5.5,6.0")
    res = np.loadtxt(txt, dtype=dtype, delimiter=",", max_rows=2)
    assert_equal(res.dtype, dtype)
    assert_equal(res, np.array([["1.5", "2.5"], ["3.0", "4.0"]], dtype=dtype))

@pytest.mark.skipif(IS_PYPY and sys.implementation.version <= (7, 3, 8),
                    reason="PyPy bug in error formatting")
@pytest.mark.parametrize("dtype", (np.dtype("f8"), np.dtype("i2")))
def test_exception_message_bad_values(dtype):
    txt = StringIO("1,2\n3,XXX\n5,6")
    msg = f"could not convert string 'XXX' to {dtype} at row 1, column 2"
    with pytest.raises(ValueError, match=msg):
        np.loadtxt(txt, dtype=dtype, delimiter=",")

def test_converters_negative_indices():
    txt = StringIO('1.5,2.5\n3.0,XXX\n5.5,6.0')
    conv = {-1: lambda s: np.nan if s == 'XXX' else float(s)}
    expected = np.array([[1.5, 2.5], [3.0, np.nan], [5.5, 6.0]])
    res = np.loadtxt(
        txt, dtype=np.float64, delimiter=",", converters=conv, encoding=None
    )
    assert_equal(res, expected)

def test_converters_negative_indices_with_usecols():
    txt = StringIO('1.5,2.5,3.5\n3.0,4.0,XXX\n5.5,6.0,7.5\n')
    conv = {-1: lambda s: np.nan if s == 'XXX' else float(s)}
    expected = np.array([[1.5, 3.5], [3.0, np.nan], [5.5, 7.5]])
    res = np.loadtxt(
        txt,
        dtype=np.float64,
        delimiter=",",
        converters=conv,
        usecols=[0, -1],
        encoding=None,
    )
    assert_equal(res, expected)

    # Second test with variable number of rows:
    res = np.loadtxt(StringIO('''0,1,2\n0,1,2,3,4'''), delimiter=",",
                     usecols=[0, -1], converters={-1: (lambda x: -1)})
    assert_array_equal(res, [[0, -1], [0, -1]])

def test_ragged_usecols():
    # usecols, and negative ones, work even with varying number of columns.
    txt = StringIO("0,0,XXX\n0,XXX,0,XXX\n0,XXX,XXX,0,XXX\n")
    expected = np.array([[0, 0], [0, 0], [0, 0]])
    res = np.loadtxt(txt, dtype=float, delimiter=",", usecols=[0, -2])
    assert_equal(res, expected)

    txt = StringIO("0,0,XXX\n0\n0,XXX,XXX,0,XXX\n")
    with pytest.raises(ValueError,
                match="invalid column index -2 at row 1 with 2 columns"):
        # There is no -2 column in the second row:
        np.loadtxt(txt, dtype=float, delimiter=",", usecols=[0, -2])

def test_empty_usecols():
    txt = StringIO("0,0,XXX\n0,XXX,0,XXX\n0,XXX,XXX,0,XXX\n")
    res = np.loadtxt(txt, dtype=np.dtype([]), delimiter=",", usecols=[])
    assert res.shape == (3,)
    assert res.dtype == np.dtype([])

@pytest.mark.parametrize("c1", ["a", "の", "🫕"])
@pytest.mark.parametrize("c2", ["a", "の", "🫕"])
def test_large_unicode_characters(c1, c2):
    # c1 and c2 span ascii, 16bit and 32bit range.
    txt = StringIO(f"a,{c1},c,1.0\ne,{c2},2.0,g")
    res = np.loadtxt(txt, dtype=np.dtype('U12'), delimiter=",")
    expected = np.array(
        [f"a,{c1},c,1.0".split(","), f"e,{c2},2.0,g".split(",")],
        dtype=np.dtype('U12')
    )
    assert_equal(res, expected)

def test_unicode_with_converter():
    txt = StringIO("cat,dog\nαβγ,δεζ\nabc,def\n")
    conv = {0: lambda s: s.upper()}
    res = np.loadtxt(
        txt,
        dtype=np.dtype("U12"),
        converters=conv,
        delimiter=",",
        encoding=None
    )
    expected = np.array([['CAT', 'dog'], ['ΑΒΓ', 'δεζ'], ['ABC', 'def']])
    assert_equal(res, expected)

def test_converter_with_structured_dtype():
    txt = StringIO('1.5,2.5,Abc\n3.0,4.0,dEf\n5.5,6.0,ghI\n')
    dt = np.dtype([('m', np.int32), ('r', np.float32), ('code', 'U8')])
    conv = {0: lambda s: int(10*float(s)), -1: lambda s: s.upper()}
    res = np.loadtxt(txt, dtype=dt, delimiter=",", converters=conv)
    expected = np.array(
        [(15, 2.5, 'ABC'), (30, 4.0, 'DEF'), (55, 6.0, 'GHI')], dtype=dt
    )
    assert_equal(res, expected)

def test_converter_with_unicode_dtype():
    """
    With the default 'bytes' encoding, tokens are encoded prior to being
    passed to the converter. This means that the output of the converter may
    be bytes instead of unicode as expected by `read_rows`.

    This test checks that outputs from the above scenario are properly decoded
    prior to parsing by `read_rows`.
    """
    txt = StringIO('abc,def\nrst,xyz')
    conv = bytes.upper
    res = np.loadtxt(
            txt, dtype=np.dtype("U3"), converters=conv, delimiter=",")
    expected = np.array([['ABC', 'DEF'], ['RST', 'XYZ']])
    assert_equal(res, expected)

def test_read_huge_row():
    row = "1.5, 2.5," * 50000
    row = row[:-1] + "\n"
    txt = StringIO(row * 2)
    res = np.loadtxt(txt, delimiter=",", dtype=float)
    assert_equal(res, np.tile([1.5, 2.5], (2, 50000)))

@pytest.mark.parametrize("dtype", "edfgFDG")
def test_huge_float(dtype):
    # Covers a non-optimized path that is rarely taken:
    field = "0" * 1000 + ".123456789"
    dtype = np.dtype(dtype)
    value = np.loadtxt([field], dtype=dtype)[()]
    assert value == dtype.type("0.123456789")

@pytest.mark.parametrize(
    ("given_dtype", "expected_dtype"),
    [
        ("S", np.dtype("S5")),
        ("U", np.dtype("U5")),
    ],
)
def test_string_no_length_given(given_dtype, expected_dtype):
    """
    The given dtype is just 'S' or 'U' with no length. In these cases, the
    length of the resulting dtype is determined by the longest string found
    in the file.
    """
    txt = StringIO("AAA,5-1\nBBBBB,0-3\nC,4-9\n")
    res = np.loadtxt(txt, dtype=given_dtype, delimiter=",")
    expected = np.array(
        [['AAA', '5-1'], ['BBBBB', '0-3'], ['C', '4-9']], dtype=expected_dtype
    )
    assert_equal(res, expected)
    assert_equal(res.dtype, expected_dtype)

def test_float_conversion():
    """
    Some tests that the conversion to float64 works as accurately as the
    Python built-in `float` function. In a naive version of the float parser,
    these strings resulted in values that were off by an ULP or two.
    """
    strings = [
        '0.9999999999999999',
        '9876543210.123456',
        '5.43215432154321e+300',
        '0.901',
        '0.333',
    ]
    txt = StringIO('\n'.join(strings))
    res = np.loadtxt(txt)
    expected = np.array([float(s) for s in strings])
    assert_equal(res, expected)

def test_bool():
    # Simple test for bool via integer
    txt = StringIO("1, 0\n10, -1")
    res = np.loadtxt(txt, dtype=bool, delimiter=",")
    assert res.dtype == bool
    assert_array_equal(res, [[True, False], [True, True]])
    # Make sure we use only 1 and 0 on the byte level:
    assert_array_equal(res.view(np.uint8), [[1, 0], [1, 1]])

@pytest.mark.skipif(IS_PYPY and sys.implementation.version <= (7, 3, 8),
                    reason="PyPy bug in error formatting")
@pytest.mark.parametrize("dtype", np.typecodes["AllInteger"])
def test_integer_signs(dtype):
    dtype = np.dtype(dtype)
    assert np.loadtxt(["+2"], dtype=dtype) == 2
    if dtype.kind == "u":
        with pytest.raises(ValueError):
            np.loadtxt(["-1\n"], dtype=dtype)
    else:
        assert np.loadtxt(["-2\n"], dtype=dtype) == -2

    for sign in ["++", "+-", "--", "-+"]:
        with pytest.raises(ValueError):
            np.loadtxt([f"{sign}2\n"], dtype=dtype)

@pytest.mark.skipif(IS_PYPY and sys.implementation.version <= (7, 3, 8),
                    reason="PyPy bug in error formatting")
@pytest.mark.parametrize("dtype", np.typecodes["AllInteger"])
def test_implicit_cast_float_to_int_fails(dtype):
    txt = StringIO("1.0, 2.1, 3.7\n4, 5, 6")
    with pytest.raises(ValueError):
        np.loadtxt(txt, dtype=dtype, delimiter=",")

@pytest.mark.parametrize("dtype", (np.complex64, np.complex128))
@pytest.mark.parametrize("with_parens", (False, True))
def test_complex_parsing(dtype, with_parens):
    s = "(1.0-2.5j),3.75,(7+-5.0j)\n(4),(-19e2j),(0)"
    if not with_parens:
        s = s.replace("(", "").replace(")", "")

    res = np.loadtxt(StringIO(s), dtype=dtype, delimiter=",")
    expected = np.array(
        [[1.0-2.5j, 3.75, 7-5j], [4.0, -1900j, 0]], dtype=dtype
    )
    assert_equal(res, expected)

def test_read_from_generator():
    def gen():
        for i in range(4):
            yield f"{i},{2*i},{i**2}"

    res = np.loadtxt(gen(), dtype=int, delimiter=",")
    expected = np.array([[0, 0, 0], [1, 2, 1], [2, 4, 4], [3, 6, 9]])
    assert_equal(res, expected)

def test_read_from_generator_multitype():
    def gen():
        for i in range(3):
            yield f"{i} {i / 4}"

    res = np.loadtxt(gen(), dtype="i, d", delimiter=" ")
    expected = np.array([(0, 0.0), (1, 0.25), (2, 0.5)], dtype="i, d")
    assert_equal(res, expected)

def test_read_from_bad_generator():
    def gen():
        for entry in ["1,2", b"3, 5", 12738]:
            yield entry

    with pytest.raises(
            TypeError, match=r"non-string returned while reading data"):
        np.loadtxt(gen(), dtype="i, i", delimiter=",")

@pytest.mark.skipif(not HAS_REFCOUNT, reason="Python lacks refcounts")
def test_object_cleanup_on_read_error():
    sentinel = object()
    already_read = 0

    def conv(x):
        nonlocal already_read
        if already_read > 4999:
            raise ValueError("failed half-way through!")
        already_read += 1
        return sentinel

    txt = StringIO("x\n" * 10000)

    with pytest.raises(ValueError, match="at row 5000, column 1"):
        np.loadtxt(txt, dtype=object, converters={0: conv})

    assert sys.getrefcount(sentinel) == 2

@pytest.mark.skipif(IS_PYPY and sys.implementation.version <= (7, 3, 8),
                    reason="PyPy bug in error formatting")
def test_character_not_bytes_compatible():
    """Test exception when a character cannot be encoded as 'S'."""
    data = StringIO("–")  # == \u2013
    with pytest.raises(ValueError):
        np.loadtxt(data, dtype="S5")

@pytest.mark.parametrize("conv", (0, [float], ""))
def test_invalid_converter(conv):
    msg = (
        "converters must be a dictionary mapping columns to converter "
        "functions or a single callable."
    )
    with pytest.raises(TypeError, match=msg):
        np.loadtxt(StringIO("1 2\n3 4"), converters=conv)

@pytest.mark.skipif(IS_PYPY and sys.implementation.version <= (7, 3, 8),
                    reason="PyPy bug in error formatting")
def test_converters_dict_raises_non_integer_key():
    with pytest.raises(TypeError, match="keys of the converters dict"):
        np.loadtxt(StringIO("1 2\n3 4"), converters={"a": int})
    with pytest.raises(TypeError, match="keys of the converters dict"):
        np.loadtxt(StringIO("1 2\n3 4"), converters={"a": int}, usecols=0)

@pytest.mark.parametrize("bad_col_ind", (3, -3))
def test_converters_dict_raises_non_col_key(bad_col_ind):
    data = StringIO("1 2\n3 4")
    with pytest.raises(ValueError, match="converter specified for column"):
        np.loadtxt(data, converters={bad_col_ind: int})

def test_converters_dict_raises_val_not_callable():
    with pytest.raises(TypeError,
                match="values of the converters dictionary must be callable"):
        np.loadtxt(StringIO("1 2\n3 4"), converters={0: 1})

@pytest.mark.parametrize("q", ('"', "'", "`"))
def test_quoted_field(q):
    txt = StringIO(
        f"{q}alpha, x{q}, 2.5\n{q}beta, y{q}, 4.5\n{q}gamma, z{q}, 5.0\n"
    )
    dtype = np.dtype([('f0', 'U8'), ('f1', np.float64)])
    expected = np.array(
        [("alpha, x", 2.5), ("beta, y", 4.5), ("gamma, z", 5.0)], dtype=dtype
    )

    res = np.loadtxt(txt, dtype=dtype, delimiter=",", quotechar=q)
    assert_array_equal(res, expected)

def test_quote_support_default():
    """Support for quoted fields is disabled by default."""
    txt = StringIO('"lat,long", 45, 30\n')
    dtype = np.dtype([('f0', 'U24'), ('f1', np.float64), ('f2', np.float64)])

    with pytest.raises(ValueError, match="the number of columns changed"):
        np.loadtxt(txt, dtype=dtype, delimiter=",")

    # Enable quoting support with non-None value for quotechar param
    txt.seek(0)
    expected = np.array([("lat,long", 45., 30.)], dtype=dtype)

    res = np.loadtxt(txt, dtype=dtype, delimiter=",", quotechar='"')
    assert_array_equal(res, expected)

@pytest.mark.skipif(IS_PYPY and sys.implementation.version <= (7, 3, 8),
                    reason="PyPy bug in error formatting")
def test_quotechar_multichar_error():
    txt = StringIO("1,2\n3,4")
    msg = r".*must be a single unicode character or None"
    with pytest.raises(TypeError, match=msg):
        np.loadtxt(txt, delimiter=",", quotechar="''")

def test_comment_multichar_error_with_quote():
    txt = StringIO("1,2\n3,4")
    msg = (
        "when multiple comments or a multi-character comment is given, "
        "quotes are not supported."
    )
    with pytest.raises(ValueError, match=msg):
        np.loadtxt(txt, delimiter=",", comments="123", quotechar='"')
    with pytest.raises(ValueError, match=msg):
        np.loadtxt(txt, delimiter=",", comments=["#", "%"], quotechar='"')

    # A single character string in a tuple is unpacked though:
    res = np.loadtxt(txt, delimiter=",", comments=("#",), quotechar="'")
    assert_equal(res, [[1, 2], [3, 4]])

def test_structured_dtype_with_quotes():
    data = StringIO(
        (
            "1000;2.4;'alpha';-34\n"
            "2000;3.1;'beta';29\n"
            "3500;9.9;'gamma';120\n"
            "4090;8.1;'delta';0\n"
            "5001;4.4;'epsilon';-99\n"
            "6543;7.8;'omega';-1\n"
        )
    )
    dtype = np.dtype(
        [('f0', np.uint16), ('f1', np.float64), ('f2', 'S7'), ('f3', np.int8)]
    )
    expected = np.array(
        [
            (1000, 2.4, "alpha", -34),
            (2000, 3.1, "beta", 29),
            (3500, 9.9, "gamma", 120),
            (4090, 8.1, "delta", 0),
            (5001, 4.4, "epsilon", -99),
            (6543, 7.8, "omega", -1)
        ],
        dtype=dtype
    )
    res = np.loadtxt(data, dtype=dtype, delimiter=";", quotechar="'")
    assert_array_equal(res, expected)

def test_quoted_field_is_not_empty():
    txt = StringIO('1\n\n"4"\n""')
    expected = np.array(["1", "4", ""], dtype="U1")
    res = np.loadtxt(txt, delimiter=",", dtype="U1", quotechar='"')
    assert_equal(res, expected)

def test_quoted_field_is_not_empty_nonstrict():
    # Same as test_quoted_field_is_not_empty but check that we are not strict
    # about missing closing quote (this is the `csv.reader` default also)
    txt = StringIO('1\n\n"4"\n"')
    expected = np.array(["1", "4", ""], dtype="U1")
    res = np.loadtxt(txt, delimiter=",", dtype="U1", quotechar='"')
    assert_equal(res, expected)

def test_consecutive_quotechar_escaped():
    txt = StringIO('"Hello, my name is ""Monty""!"')
    expected = np.array('Hello, my name is "Monty"!', dtype="U40")
    res = np.loadtxt(txt, dtype="U40", delimiter=",", quotechar='"')
    assert_equal(res, expected)

@pytest.mark.parametrize("data", ("", "\n\n\n", "# 1 2 3\n# 4 5 6\n"))
@pytest.mark.parametrize("ndmin", (0, 1, 2))
@pytest.mark.parametrize("usecols", [None, (1, 2, 3)])
def test_warn_on_no_data(data, ndmin, usecols):
    """Check that a UserWarning is emitted when no data is read from input."""
    if usecols is not None:
        expected_shape = (0, 3)
    elif ndmin == 2:
        expected_shape = (0, 1)  # guess a single column?!
    else:
        expected_shape = (0,)

    txt = StringIO(data)
    with pytest.warns(UserWarning, match="input contained no data"):
        res = np.loadtxt(txt, ndmin=ndmin, usecols=usecols)
    assert res.shape == expected_shape

    with NamedTemporaryFile(mode="w") as fh:
        fh.write(data)
        fh.seek(0)
        with pytest.warns(UserWarning, match="input contained no data"):
            res = np.loadtxt(txt, ndmin=ndmin, usecols=usecols)
        assert res.shape == expected_shape

@pytest.mark.parametrize("skiprows", (2, 3))
def test_warn_on_skipped_data(skiprows):
    data = "1 2 3\n4 5 6"
    txt = StringIO(data)
    with pytest.warns(UserWarning, match="input contained no data"):
        np.loadtxt(txt, skiprows=skiprows)

@pytest.mark.parametrize(["dtype", "value"], [
        ("i2", 0x0001), ("u2", 0x0001),
        ("i4", 0x00010203), ("u4", 0x00010203),
        ("i8", 0x0001020304050607), ("u8", 0x0001020304050607),
        # The following values are constructed to lead to unique bytes:
        ("float16", 3.07e-05),
        ("float32", 9.2557e-41), ("complex64", 9.2557e-41+2.8622554e-29j),
        ("float64", -1.758571353180402e-24),
        # Here and below, the repr side-steps a small loss of precision in
        # complex `str` in PyPy (which is probably fine, as repr works):
        ("complex128", repr(5.406409232372729e-29-1.758571353180402e-24j)),
        # Use integer values that fit into double.  Everything else leads to
        # problems due to longdoubles going via double and decimal strings
        # causing rounding errors.
        ("longdouble", 0x01020304050607),
        ("clongdouble", repr(0x01020304050607 + (0x00121314151617 * 1j))),
        ("U2", "\U00010203\U000a0b0c")])
@pytest.mark.parametrize("swap", [True, False])
def test_byteswapping_and_unaligned(dtype, value, swap):
    # Try to create "interesting" values within the valid unicode range:
    dtype = np.dtype(dtype)
    data = [f"x,{value}\n"]  # repr as PyPy `str` truncates some
    if swap:
        dtype = dtype.newbyteorder()
    full_dt = np.dtype([("a", "S1"), ("b", dtype)], align=False)
    # The above ensures that the interesting "b" field is unaligned:
    assert full_dt.fields["b"][1] == 1
    res = np.loadtxt(data, dtype=full_dt, delimiter=",", encoding=None,
                     max_rows=1)  # max-rows prevents over-allocation
    assert res["b"] == dtype.type(value)

@pytest.mark.parametrize("dtype",
        np.typecodes["AllInteger"] + "efdFD" + "?")
def test_unicode_whitespace_stripping(dtype):
    # Test that all numeric types (and bool) strip whitespace correctly
    # \u202F is a narrow no-break space, `\n` is just a whitespace if quoted.
    # Currently, skip float128 as it did not always support this and has no
    # "custom" parsing:
    txt = StringIO(' 3 ,"\u202F2\n"')
    res = np.loadtxt(txt, dtype=dtype, delimiter=",", quotechar='"')
    assert_array_equal(res, np.array([3, 2]).astype(dtype))

@pytest.mark.parametrize("dtype", "FD")
def test_unicode_whitespace_stripping_complex(dtype):
    # Complex has a few extra cases since it has two components and
    # parentheses
    line = " 1 , 2+3j , ( 4+5j ), ( 6+-7j )  , 8j , ( 9j ) \n"
    data = [line, line.replace(" ", "\u202F")]
    res = np.loadtxt(data, dtype=dtype, delimiter=',')
    assert_array_equal(res, np.array([[1, 2+3j, 4+5j, 6-7j, 8j, 9j]] * 2))

@pytest.mark.skipif(IS_PYPY and sys.implementation.version <= (7, 3, 8),
                    reason="PyPy bug in error formatting")
@pytest.mark.parametrize("dtype", "FD")
@pytest.mark.parametrize("field",
        ["1 +2j", "1+ 2j", "1+2 j", "1+-+3", "(1j", "(1", "(1+2j", "1+2j)"])
def test_bad_complex(dtype, field):
    with pytest.raises(ValueError):
        np.loadtxt([field + "\n"], dtype=dtype, delimiter=",")

@pytest.mark.skipif(IS_PYPY and sys.implementation.version <= (7, 3, 8),
                    reason="PyPy bug in error formatting")
@pytest.mark.parametrize("dtype",
            np.typecodes["AllInteger"] + "efgdFDG" + "?")
def test_nul_character_error(dtype):
    # Test that a \0 character is correctly recognized as an error even if
    # what comes before is valid (not everything gets parsed internally).
    if dtype.lower() == "g":
        pytest.xfail("longdouble/clongdouble assignment may misbehave.")
    with pytest.raises(ValueError):
        np.loadtxt(["1\000"], dtype=dtype, delimiter=",", quotechar='"')

@pytest.mark.skipif(IS_PYPY and sys.implementation.version <= (7, 3, 8),
                    reason="PyPy bug in error formatting")
@pytest.mark.parametrize("dtype",
        np.typecodes["AllInteger"] + "efgdFDG" + "?")
def test_no_thousands_support(dtype):
    # Mainly to document behaviour, Python supports thousands like 1_1.
    # (e and G may end up using different conversion and support it, this is
    # a bug but happens...)
    if dtype == "e":
        pytest.skip("half assignment currently uses Python float converter")
    if dtype in "eG":
        pytest.xfail("clongdouble assignment is buggy (uses `complex`?).")

    assert int("1_1") == float("1_1") == complex("1_1") == 11
    with pytest.raises(ValueError):
        np.loadtxt(["1_1\n"], dtype=dtype)

@pytest.mark.parametrize("data", [
    ["1,2\n", "2\n,3\n"],
    ["1,2\n", "2\r,3\n"]])
def test_bad_newline_in_iterator(data):
    # In NumPy <=1.22 this was accepted, because newlines were completely
    # ignored when the input was an iterable.  This could be changed, but right
    # now, we raise an error.
    msg = "Found an unquoted embedded newline within a single line"
    with pytest.raises(ValueError, match=msg):
        np.loadtxt(data, delimiter=",")

@pytest.mark.parametrize("data", [
    ["1,2\n", "2,3\r\n"],  # a universal newline
    ["1,2\n", "'2\n',3\n"],  # a quoted newline
    ["1,2\n", "'2\r',3\n"],
    ["1,2\n", "'2\r\n',3\n"],
])
def test_good_newline_in_iterator(data):
    # The quoted newlines will be untransformed here, but are just whitespace.
    res = np.loadtxt(data, delimiter=",", quotechar="'")
    assert_array_equal(res, [[1., 2.], [2., 3.]])

@pytest.mark.parametrize("newline", ["\n", "\r", "\r\n"])
def test_universal_newlines_quoted(newline):
    # Check that universal newline support within the tokenizer is not applied
    # to quoted fields.  (note that lines must end in newline or quoted
    # fields will not include a newline at all)
    data = ['1,"2\n"\n', '3,"4\n', '1"\n']
    data = [row.replace("\n", newline) for row in data]
    res = np.loadtxt(data, dtype=object, delimiter=",", quotechar='"')
    assert_array_equal(res, [['1', f'2{newline}'], ['3', f'4{newline}1']])

def test_null_character():
    # Basic tests to check that the NUL character is not special:
    res = np.loadtxt(["1\0002\0003\n", "4\0005\0006"], delimiter="\000")
    assert_array_equal(res, [[1, 2, 3], [4, 5, 6]])

    # Also not as part of a field (avoid unicode/arrays as unicode strips \0)
    res = np.loadtxt(["1\000,2\000,3\n", "4\000,5\000,6"],
                     delimiter=",", dtype=object)
    assert res.tolist() == [["1\000", "2\000", "3"], ["4\000", "5\000", "6"]]

def test_iterator_fails_getting_next_line():
    class BadSequence:
        def __len__(self):
            return 100

        def __getitem__(self, item):
            if item == 50:
                raise RuntimeError("Bad things happened!")
            return f"{item}, {item+1}"

    with pytest.raises(RuntimeError, match="Bad things happened!"):
        np.loadtxt(BadSequence(), dtype=int, delimiter=",")

class TestCReaderUnitTests:
    # These are internal tests for path that should not be possible to hit
    # unless things go very very wrong somewhere.
    def test_not_an_filelike(self):
        with pytest.raises(AttributeError, match=".*read"):
            np.core._multiarray_umath._load_from_filelike(
                object(), dtype=np.dtype("i"), filelike=True)

    def test_filelike_read_fails(self):
        # Can only be reached if loadtxt opens the file, so it is hard to do
        # via the public interface (although maybe not impossible considering
        # the current "DataClass" backing).
        class BadFileLike:
            counter = 0

            def read(self, size):
                self.counter += 1
                if self.counter > 20:
                    raise RuntimeError("Bad bad bad!")
                return "1,2,3\n"

        with pytest.raises(RuntimeError, match="Bad bad bad!"):
            np.core._multiarray_umath._load_from_filelike(
                BadFileLike(), dtype=np.dtype("i"), filelike=True)

    def test_filelike_bad_read(self):
        # Can only be reached if loadtxt opens the file, so it is hard to do
        # via the public interface (although maybe not impossible considering
        # the current "DataClass" backing).

        class BadFileLike:
            counter = 0

            def read(self, size):
                return 1234  # not a string!

        with pytest.raises(TypeError,
                    match="non-string returned while reading data"):
            np.core._multiarray_umath._load_from_filelike(
                BadFileLike(), dtype=np.dtype("i"), filelike=True)

    def test_not_an_iter(self):
        with pytest.raises(TypeError,
                    match="error reading from object, expected an iterable"):
            np.core._multiarray_umath._load_from_filelike(
                object(), dtype=np.dtype("i"), filelike=False)

    def test_bad_type(self):
        with pytest.raises(TypeError, match="internal error: dtype must"):
            np.core._multiarray_umath._load_from_filelike(
                object(), dtype="i", filelike=False)

    def test_bad_encoding(self):
        with pytest.raises(TypeError, match="encoding must be a unicode"):
            np.core._multiarray_umath._load_from_filelike(
                object(), dtype=np.dtype("i"), filelike=False, encoding=123)

    @pytest.mark.parametrize("newline", ["\r", "\n", "\r\n"])
    def test_manual_universal_newlines(self, newline):
        # This is currently not available to users, because we should always
        # open files with universal newlines enabled `newlines=None`.
        # (And reading from an iterator uses slightly different code paths.)
        # We have no real support for `newline="\r"` or `newline="\n" as the
        # user cannot specify those options.
        data = StringIO('0\n1\n"2\n"\n3\n4 #\n'.replace("\n", newline),
                        newline="")

        res = np.core._multiarray_umath._load_from_filelike(
            data, dtype=np.dtype("U10"), filelike=True,
            quote='"', comment="#", skiplines=1)
        assert_array_equal(res[:, 0], ["1", f"2{newline}", "3", "4 "])

def test_delimiter_comment_collision_raises():
    with pytest.raises(TypeError, match=".*control characters.*incompatible"):
        np.loadtxt(StringIO("1, 2, 3"), delimiter=",", comments=",")

def test_delimiter_quotechar_collision_raises():
    with pytest.raises(TypeError, match=".*control characters.*incompatible"):
        np.loadtxt(StringIO("1, 2, 3"), delimiter=",", quotechar=",")

def test_comment_quotechar_collision_raises():
    with pytest.raises(TypeError, match=".*control characters.*incompatible"):
        np.loadtxt(StringIO("1 2 3"), comments="#", quotechar="#")

def test_delimiter_and_multiple_comments_collision_raises():
    with pytest.raises(
        TypeError, match="Comment characters.*cannot include the delimiter"
    ):
        np.loadtxt(StringIO("1, 2, 3"), delimiter=",", comments=["#", ","])

@pytest.mark.parametrize(
    "ws",
    (
        " ",  # space
        "\t",  # tab
        "\u2003",  # em
        "\u00A0",  # non-break
        "\u3000",  # ideographic space
    )
)
def test_collision_with_default_delimiter_raises(ws):
    with pytest.raises(TypeError, match=".*control characters.*incompatible"):
        np.loadtxt(StringIO(f"1{ws}2{ws}3\n4{ws}5{ws}6\n"), comments=ws)
    with pytest.raises(TypeError, match=".*control characters.*incompatible"):
        np.loadtxt(StringIO(f"1{ws}2{ws}3\n4{ws}5{ws}6\n"), quotechar=ws)

@pytest.mark.parametrize("nl", ("\n", "\r"))
def test_control_character_newline_raises(nl):
    txt = StringIO(f"1{nl}2{nl}3{nl}{nl}4{nl}5{nl}6{nl}{nl}")
    msg = "control character.*cannot be a newline"
    with pytest.raises(TypeError, match=msg):
        np.loadtxt(txt, delimiter=nl)
    with pytest.raises(TypeError, match=msg):
        np.loadtxt(txt, comments=nl)
    with pytest.raises(TypeError, match=msg):
        np.loadtxt(txt, quotechar=nl)

@pytest.mark.parametrize(
    ("generic_data", "long_datum", "unitless_dtype", "expected_dtype"),
    [
        ("2012-03", "2013-01-15", "M8", "M8[D]"),  # Datetimes
        ("spam-a-lot", "tis_but_a_scratch", "U", "U17"),  # str
    ],
)
@pytest.mark.parametrize("nrows", (10, 50000, 60000))  # lt, eq, gt chunksize
def test_parametric_unit_discovery(
    generic_data, long_datum, unitless_dtype, expected_dtype, nrows
):
    """Check that the correct unit (e.g. month, day, second) is discovered from
    the data when a user specifies a unitless datetime."""
    # Unit should be "D" (days) due to last entry
    data = [generic_data] * 50000 + [long_datum]
    expected = np.array(data, dtype=expected_dtype)

    # file-like path
    txt = StringIO("\n".join(data))
    a = np.loadtxt(txt, dtype=unitless_dtype)
    assert a.dtype == expected.dtype
    assert_equal(a, expected)

    # file-obj path
    fd, fname = mkstemp()
    with open(fname, "w") as fh:
        fh.write("\n".join(data))
    a = np.loadtxt(fname, dtype=unitless_dtype)
    assert a.dtype == expected.dtype
    assert_equal(a, expected)

def test_str_dtype_unit_discovery_with_converter():
    data = ["spam-a-lot"] * 60000 + ["XXXtis_but_a_scratch"]
    expected = np.array(
        ["spam-a-lot"] * 60000 + ["tis_but_a_scratch"], dtype="U17"
    )
    conv = lambda s: s.strip("XXX")

    # file-like path
    txt = StringIO("\n".join(data))
    a = np.loadtxt(txt, dtype="U", converters=conv, encoding=None)
    assert a.dtype == expected.dtype
    assert_equal(a, expected)

    # file-obj path
    fd, fname = mkstemp()
    with open(fname, "w") as fh:
        fh.write("\n".join(data))
    a = np.loadtxt(fname, dtype="U", converters=conv, encoding=None)
    assert a.dtype == expected.dtype
    assert_equal(a, expected)

@pytest.mark.skipif(IS_PYPY and sys.implementation.version <= (7, 3, 8),
                    reason="PyPy bug in error formatting")
def test_control_character_empty():
    with pytest.raises(TypeError, match="Text reading control character must"):
        np.loadtxt(StringIO("1 2 3"), delimiter="")
    with pytest.raises(TypeError, match="Text reading control character must"):
        np.loadtxt(StringIO("1 2 3"), quotechar="")
    with pytest.raises(ValueError, match="comments cannot be an empty string"):
        np.loadtxt(StringIO("1 2 3"), comments="")
    with pytest.raises(ValueError, match="comments cannot be an empty string"):
        np.loadtxt(StringIO("1 2 3"), comments=["#", ""])

def test_control_characters_as_bytes():
    """Byte control characters (comments, delimiter) are supported."""
    a = np.loadtxt(StringIO("#header\n1,2,3"), comments=b"#", delimiter=b",")
    assert_equal(a, [1, 2, 3])