# -*- coding: utf-8 -*-

"""
Script defining all the functions that are specific to the MNIST dataset
"""

from __future__ import annotations

import os
from time import perf_counter
from typing import Union, Optional

import numpy as np
import matplotlib.pyplot as plt

import utils
from utils import (
    cast,
    check_dtype,
    get_dtype_of_array,
    get_range_of_array,
    vector_to_categorical,
    categorical_to_vector,
    list_to_string,
    display_class_distributions,
    _download_data,
    _validate_selected_classes
)

from core import train_test_split

##############################################################################

def _download_raw_MNIST_dataset() -> str:
    r"""
    Automatically downloads the raw MNIST data (as a single file), and saves
    it to the following location on your disk :
        - on Windows : "C:\Users\YourUsername\.Custom-NN\datasets\MNIST\raw_MNIST_data.npz"
        - on Linux   : "/home/YourUsername/.Custom-NN/datasets/MNIST/raw_MNIST_data.npz"

    This function is basically a wrapper of the `_download_data` function of
    the "utils.py" script

    Naturally, if this is the very first time you call this function, you'll
    need to have an internet connection !

    The downloaded file has a size of about 11 MB, and is retrieved from the
    following URL :
    https://storage.googleapis.com/tensorflow/tf-keras-datasets/mnist.npz

    Also, the absolute path of the downloaded data is returned

    Sidenote
    --------
    This function is basically equivalent to the `tensorflow.keras.datasets.mnist.load_data`
    method. The only real difference is that the downloaded data has a different
    location on your disk. The reason why the previous `mnist.load_data`
    method isn't used is simply because we do NOT want to import TensorFlow !
    Indeed, in my opinion, it would be kind of awkward to import TensorFlow in
    a project aiming to implement a Deep Learning model *from scratch* !
    Note that, if you use the `mnist.load_data` method of the TensorFlow module,
    the raw MNIST data will be saved at the following location on your disk :
        - on Windows : "C:\Users\YourUsername\.keras\datasets\mnist.npz"
        - on Linux   : "/home/YourUsername/.keras/datasets/mnist.npz"
    """
    # ---------------------------------------------------------------------- #

    # the raw MNIST data will be downloaded from this URL (by default)

    data_URL = "https://storage.googleapis.com/tensorflow/tf-keras-datasets/mnist.npz"

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

    # creating the default folder that will contain the raw MNIST data (if it
    # doesn't already exist)
    default_data_directory = os.path.join(
        os.path.expanduser("~"),
        ".Custom-NN",
        "datasets",
        "MNIST"
    )
    os.makedirs(default_data_directory, exist_ok=True)

    # here, `default_data_filename` must have a "*.npz" extension
    default_data_filename = "raw_MNIST_data.npz"
    assert os.path.splitext(default_data_filename)[1] == ".npz"

    # defining the absolute path of the downloaded file
    path_of_downloaded_data = os.path.join(
        default_data_directory,
        default_data_filename
    )

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

    # default SHA-256 hash value of the raw MNIST data

    hash_value = "731c5ac602752760c8e48fbffcf8c3b850d9dc2a2aedcf2cc48468fc17b673d1"

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

    # actually downloading the raw MNIST data (if it doesn't already exist
    # on your disk)

    _download_data(
        data_URL,
        path_of_downloaded_data,
        hash_value=hash_value
    )

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

    return path_of_downloaded_data

##############################################################################

def _validate_raw_MNIST_dataset(
        raw_X_train: np.ndarray,
        raw_y_train: np.ndarray,
        raw_X_test:  np.ndarray,
        raw_y_test:  np.ndarray
    ) -> None:
    """
    Checks if the specified raw MNIST data is valid or not. By design, the
    arguments `raw_X_train`, `raw_y_train`, `raw_X_test` and `raw_y_test` are
    meant to be the outputs of the `load_raw_MNIST_dataset_from_disk` function
    (of the "mnist_dataset.py" script)
    """
    assert isinstance(raw_X_train, np.ndarray)
    assert raw_X_train.shape == (60000, 28, 28)
    check_dtype(raw_X_train, np.uint8)

    assert isinstance(raw_y_train, np.ndarray)
    assert raw_y_train.shape == (60000, )
    check_dtype(raw_y_train, np.uint8)

    assert isinstance(raw_X_test, np.ndarray)
    assert raw_X_test.shape == (10000, 28, 28)
    check_dtype(raw_X_test, np.uint8)

    assert isinstance(raw_y_test, np.ndarray)
    assert raw_y_test.shape == (10000, )
    check_dtype(raw_y_test, np.uint8)

    DEFAULT_NB_CLASSES = 10
    expected_classes = np.arange(DEFAULT_NB_CLASSES).astype(np.uint8)
    assert np.allclose(np.unique(raw_y_train), expected_classes)
    assert np.allclose(np.unique(raw_y_test),  expected_classes)

##############################################################################

def load_raw_MNIST_dataset_from_disk(
        *,
        verbose: bool = False
    ) -> tuple[np.ndarray]:
    """
    Loads the raw MNIST data from the disk
    """
    assert isinstance(verbose, bool)

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

    # downloading the raw MNIST data to the (default) location `path_of_downloaded_data`
    # (cf. the "_download_raw_MNIST_dataset" function of this script), if it
    # hasn't already been done
    path_of_downloaded_data = _download_raw_MNIST_dataset()

    t_beginning_data_loading = perf_counter()

    # actually loading the raw MNIST data from the disk
    with np.load(path_of_downloaded_data, allow_pickle=True) as RAW_MNIST_DATA:
        raw_X_train = RAW_MNIST_DATA["x_train"]
        raw_y_train = RAW_MNIST_DATA["y_train"]
        raw_X_test  = RAW_MNIST_DATA["x_test"]
        raw_y_test  = RAW_MNIST_DATA["y_test"]

    # just to be *100% sure* the loaded data is valid
    _validate_raw_MNIST_dataset(
        raw_X_train,
        raw_y_train,
        raw_X_test,
        raw_y_test
    )

    t_end_data_loading = perf_counter()
    duration_data_loading = t_end_data_loading - t_beginning_data_loading

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

    if verbose:
        print("\nShapes of the raw MNIST data :")
        print(f"    - X_train : {raw_X_train.shape}")
        print(f"    - y_train : {raw_y_train.shape}")
        print(f"    - X_test  : {raw_X_test.shape}")
        print(f"    - y_test  : {raw_y_test.shape}")

        print("\nTypes of the raw MNIST data :")
        print(f"    - X_train : {get_dtype_of_array(raw_X_train)}")
        print(f"    - y_train : {get_dtype_of_array(raw_y_train)}")
        print(f"    - X_test  : {get_dtype_of_array(raw_X_test)}")
        print(f"    - y_test  : {get_dtype_of_array(raw_y_test)}")

        precision_of_printed_info = 3
        print("\nRanges of the raw MNIST data :")
        print(f"    - X_train : {get_range_of_array(raw_X_train, precision=precision_of_printed_info)}")
        print(f"    - y_train : {get_range_of_array(raw_y_train, precision=precision_of_printed_info)}")
        print(f"    - X_test  : {get_range_of_array(raw_X_test,  precision=precision_of_printed_info)}")
        print(f"    - y_test  : {get_range_of_array(raw_y_test,  precision=precision_of_printed_info)}")

        # displaying the class distribution of the raw MNIST data
        dict_of_label_vectors = {
            "y_train" : raw_y_train,
            "y_test"  : raw_y_test
        }
        display_class_distributions(
            dict_of_label_vectors,
            precision=2
        )

    print(f"\nThe raw MNIST dataset was successfully loaded from the disk. Done in {duration_data_loading:.3f} seconds")

    return raw_X_train, raw_y_train, raw_X_test, raw_y_test

##############################################################################

def plot_random_images_from_raw_MNIST_dataset(
        raw_X_train: np.ndarray,
        raw_y_train: np.ndarray,
        raw_X_test:  np.ndarray,
        raw_y_test:  np.ndarray,
        *,
        seed: Optional[int] = None
    ) -> None:
    """
    Plots a random sample of each of the 10 digits (from the raw MNIST data).
    By design, the arguments `raw_X_train`, `raw_y_train`, `raw_X_test` and
    `raw_y_test` are meant to be the outputs of the `load_raw_MNIST_dataset_from_disk`
    function of this script
    """
    # ---------------------------------------------------------------------- #

    # checking the validity of the specified arguments

    _validate_raw_MNIST_dataset(
        raw_X_train,
        raw_y_train,
        raw_X_test,
        raw_y_test
    )

    assert isinstance(seed, (type(None), int))
    if seed is not None:
        assert seed >= 0

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

    nb_rows    = 2
    nb_columns = 5

    nb_classes = np.unique(raw_y_train).size # = 10
    assert nb_rows * nb_columns == nb_classes

    data_types = ["train", "test"]

    fig, ax = plt.subplots(nb_rows, nb_columns, figsize=(16, 8))
    plt.suptitle("\nRandom MNIST samples of each digit", fontsize=15)

    np.random.seed(seed)

    for image_index in range(nb_classes):
        random_data_type = data_types[np.random.randint(0, len(data_types))]
        if random_data_type == "train":
            # raw "train" data
            data   = raw_X_train
            labels = raw_y_train
        else:
            # raw "test" data
            data   = raw_X_test
            labels = raw_y_test

        possible_image_indices = np.where(labels == image_index)[0]
        random_image_index = possible_image_indices[np.random.randint(0, possible_image_indices.size)]
        random_image = data[random_image_index]
        assert random_image.ndim == 2

        row_index = image_index // nb_columns
        column_index = image_index % nb_columns
        ax[row_index, column_index].imshow(random_image, cmap="gray")

    # resetting the seed
    np.random.seed(None)

    plt.subplots_adjust(wspace=0.5, top=0.85)

    plt.show()

##############################################################################

def format_raw_MNIST_dataset(
        raw_X_train: np.ndarray,
        raw_y_train: np.ndarray,
        raw_X_test:  np.ndarray,
        raw_y_test:  np.ndarray,
        nb_train_samples: int,
        nb_val_samples:   int,
        nb_test_samples:  int,
        *,
        selected_classes: Union[str, list[int]] = "all",
        dict_of_real_class_names: Optional[dict[int, str]] = None,
        nb_shuffles: int = 20,
        seed: Optional[int] = None,
        verbose: bool = False
    ) -> tuple[Union[np.ndarray, None]]:
    """
    Formats the raw MNIST data so that it can be directly interpreted by a
    regular MLP neural network. By design, the arguments `raw_X_train`,
    `raw_y_train`, `raw_X_test` and `raw_y_test` are meant to be the outputs
    of the `load_raw_MNIST_dataset_from_disk` function of this script

    This function returns `X_train`, `y_train`, `X_val`, `y_val`, `X_test`
    and `y_test`

    If you do not want to get validation data, you can simply set the argument
    `nb_val_samples` to zero. In that case, the outputs `X_val` and `y_val`
    will both be equal to `None`

    The kwarg `selected_classes` can either be :
        - the string "all", if you want to work with all the classes (default)
        - a list/tuple/1D-array containing the specific class indices you want
          to work with (e.g. [2, 4, 7])

    The kwarg `dict_of_real_class_names`, if not set to `None`, is a dictionary
    with :
        - as its keys   : all the selected class indices (as integers)
        - as its values : the REAL names of the associated classes (as strings)
    For instance, if you set `selected_classes` to `[2, 4, 7]`, then you
    could, for instance, set `dict_of_real_class_names` to the following
    dictionary :
    dict_of_real_class_names = {
        2 : "TWO",
        4 : "FOUR",
        7 : "SEVEN"
    }
    By default, if `dict_of_real_class_names` is set to `None`, then the
    class names will simply be the string representations of the (selected)
    class indices
    """
    # ====================================================================== #

    # checking the validity of the specified arguments

    _validate_raw_MNIST_dataset(
        raw_X_train,
        raw_y_train,
        raw_X_test,
        raw_y_test
    )

    nb_classes = np.unique(raw_y_train).size # = 10

    # the 3 arguments `nb_train_samples`, `nb_val_samples` and `nb_test_samples`
    # will be re-checked a bit later in this function, once we check if the
    # `selected_classes` kwarg is equal to "all" or not (so that we can get
    # the actual number of classes in the - selected - raw MNIST data)
    assert isinstance(nb_train_samples, int)
    assert nb_train_samples > 0
    assert isinstance(nb_val_samples, int)
    assert nb_val_samples >= 0
    _return_validation_data = (nb_val_samples > 0)
    assert isinstance(nb_test_samples, int)
    assert nb_test_samples > 0

    selected_classes, _ = _validate_selected_classes(
        selected_classes,
        nb_classes,
        dict_of_real_class_names=dict_of_real_class_names
    )

    assert isinstance(nb_shuffles, int)
    assert nb_shuffles >= 0

    assert isinstance(seed, (type(None), int))
    if seed is not None:
        assert seed >= 0

    assert isinstance(verbose, bool)

    # ====================================================================== #

    t_beginning_data_formatting = perf_counter()

    # ====================================================================== #

    # only keeping the selected classes in the raw MNIST data

    if not(isinstance(selected_classes, str)):
        nb_classes = selected_classes.size

        indices_of_selected_train_classes = []
        indices_of_selected_test_classes  = []

        for selected_class in selected_classes:
            indices_of_selected_train_class = np.where(raw_y_train == selected_class)[0]
            indices_of_selected_train_classes.append(indices_of_selected_train_class)

            indices_of_selected_test_class = np.where(raw_y_test == selected_class)[0]
            indices_of_selected_test_classes.append(indices_of_selected_test_class)

        indices_of_selected_train_classes = np.hstack(tuple(indices_of_selected_train_classes))
        indices_of_selected_test_classes  = np.hstack(tuple(indices_of_selected_test_classes))

        # actually retrieving the selected classes from the raw data
        raw_X_train = raw_X_train[indices_of_selected_train_classes, :]
        raw_y_train = raw_y_train[indices_of_selected_train_classes]
        raw_X_test  = raw_X_test[indices_of_selected_test_classes, :]
        raw_y_test  = raw_y_test[indices_of_selected_test_classes]

        assert np.allclose(np.unique(raw_y_train), selected_classes)
        assert np.allclose(np.unique(raw_y_test),  selected_classes)

        print(f"\nNB : With the selected classes {list_to_string(selected_classes)}, here are the new shapes of the raw \"train\" and \"test\" data (from which the formatted data will be extracted) :")
        print(f"    - X_train : {raw_X_train.shape}")
        print(f"    - y_train : {raw_y_train.shape}")
        print(f"    - X_test  : {raw_X_test.shape}")
        print(f"    - y_test  : {raw_y_test.shape}")

    # ====================================================================== #

    # re-checking the validity of the arguments `nb_train_samples`, `nb_val_samples`
    # and `nb_test_samples` (since the raw data might have been modified if
    # all the classes weren't selected)

    # we want to have each class represented AT LEAST ONCE in the formatted
    # data, therefore the number of samples of the train, val and test datasets
    # CANNOT be strictly less than `nb_classes`
    assert nb_train_samples >= nb_classes
    if _return_validation_data:
        assert nb_val_samples >= nb_classes
    assert nb_test_samples >= nb_classes

    # NB : The validation set will be extracted from the raw training set
    total_nb_of_raw_train_samples = raw_X_train.shape[0] # = 60000 (if ALL the classes are selected)
    assert nb_train_samples + nb_val_samples <= total_nb_of_raw_train_samples

    total_nb_of_raw_test_samples = raw_X_test.shape[0] # = 10000 (if ALL the classes are selected)
    assert nb_test_samples <= total_nb_of_raw_test_samples

    # ====================================================================== #

    # step 1/6 : copying, reshaping, casting and normalizing the data

    # copying the raw data
    X_train = raw_X_train.copy()
    y_train = raw_y_train.copy()
    X_test  = raw_X_test.copy()
    y_test  = raw_y_test.copy()

    # flattening the images/samples
    nb_pixels_per_image = raw_X_train.shape[1] * raw_X_train.shape[2] # = 28 * 28 = 784
    X_train = X_train.reshape((total_nb_of_raw_train_samples, nb_pixels_per_image))
    X_test  = X_test.reshape((total_nb_of_raw_test_samples, nb_pixels_per_image))

    # casting the "uint8" data into `utils.DEFAULT_DATATYPE` data
    X_train = cast(X_train, utils.DEFAULT_DATATYPE)
    X_test  = cast(X_test,  utils.DEFAULT_DATATYPE)

    # casting the "uint8" labels into "int32" labels (optional)
    y_train = np.int_(y_train)
    y_test  = np.int_(y_test)

    # normalizing the data between 0 and 1 (by convention)
    normalizing_factor = cast(1, utils.DEFAULT_DATATYPE) / cast(255, utils.DEFAULT_DATATYPE)
    X_train *= normalizing_factor
    X_test  *= normalizing_factor

    # ====================================================================== #

    # step 2/6 : splitting the data into train/test or train/val/test
    #            subsets, such that the resulting class distributions are
    #            roughly the same as the initial raw data

    if _return_validation_data:
        # getting the formatted "train" and "val" datasets (from the raw "train" data)
        X_train, X_val, y_train, y_val = train_test_split(
            X_train,
            y_train,
            train_size=nb_train_samples,
            test_size=nb_val_samples,
            stratify=True, # keeping the same class distribution as `y_train`
            random_state=seed
        )
    else:
        # getting the formatted "train" dataset (from the raw "train" data)

        if total_nb_of_raw_train_samples - nb_train_samples >= nb_classes:
            X_train, _, y_train, _ = train_test_split(
                X_train,
                y_train,
                train_size=nb_train_samples,
                stratify=True, # keeping the same class distribution as `y_train`
                random_state=seed
            )
        else:
            # in this specific case, the `train_test_split` function breaks
            # (with the arguments used in the previous `if` block)

            np.random.seed(seed)
            train_indices = np.random.choice(
                np.arange(total_nb_of_raw_train_samples),
                size=(nb_train_samples, ),
                replace=False
            )
            np.random.seed(None) # resetting the seed

            X_train = X_train[train_indices, :].copy()
            y_train = y_train[train_indices].copy()

    assert np.unique(y_train).size == nb_classes
    if _return_validation_data:
        assert np.unique(y_val).size == nb_classes

    if seed is not None:
        # updating the seed such that the "randomness" of the test split is
        # different than the one used in the previous train (or train/val )
        # split (just in case)
        seed += 1 # the increment value is arbitrary

    # getting the formatted "test" dataset (from the raw "test" data)
    if total_nb_of_raw_test_samples - nb_test_samples >= nb_classes:
        X_test, _, y_test, _ = train_test_split(
            X_test,
            y_test,
            train_size=nb_test_samples,
            stratify=True, # keeping the same class distribution as `y_test`
            random_state=seed
        )
    else:
        # in this specific case, the `train_test_split` function breaks
        # (with the arguments used in the previous `if` block)

        np.random.seed(seed)
        test_indices = np.random.choice(
            np.arange(total_nb_of_raw_test_samples),
            size=(nb_test_samples, ),
            replace=False
        )
        np.random.seed(None) # resetting the seed

        X_test = X_test[test_indices, :].copy()
        y_test = y_test[test_indices].copy()

    assert np.unique(y_test).size == nb_classes

    if (nb_shuffles > 0) and (seed is not None):
        # updating the seed such that the "randomness" of the shuffle is
        # different than the one used in the previous splits (just in case)
        seed += 1 # the increment value is arbitrary

    # ====================================================================== #

    # step 3/6 : one-hot encoding the resulting label vectors

    y_train = vector_to_categorical(y_train, dtype=utils.DEFAULT_DATATYPE)
    if _return_validation_data:
        y_val = vector_to_categorical(y_val, dtype=utils.DEFAULT_DATATYPE)
    y_test = vector_to_categorical(y_test,  dtype=utils.DEFAULT_DATATYPE)

    # ====================================================================== #

    # step 4/6 : shuffling the data

    if nb_shuffles > 0:
        shuffled_indices_train_data = np.arange(nb_train_samples)
        if _return_validation_data:
            shuffled_indices_validation_data = np.arange(nb_val_samples)
        shuffled_indices_test_data = np.arange(nb_test_samples)

        np.random.seed(seed)
        for shuffle_index in range(nb_shuffles):
            np.random.shuffle(shuffled_indices_train_data)
            if _return_validation_data:
                np.random.shuffle(shuffled_indices_validation_data)
            np.random.shuffle(shuffled_indices_test_data)
        np.random.seed(None) # resetting the seed

        X_train = X_train[shuffled_indices_train_data, :]
        y_train = y_train[shuffled_indices_train_data, :]
        if _return_validation_data:
            X_val = X_val[shuffled_indices_validation_data, :]
            y_val = y_val[shuffled_indices_validation_data, :]
        X_test = X_test[shuffled_indices_test_data, :]
        y_test = y_test[shuffled_indices_test_data, :]

    # ====================================================================== #

    # step 5/6 : checking the shapes of the final formatted data

    assert X_train.shape == (nb_train_samples, nb_pixels_per_image)
    assert y_train.shape == (nb_train_samples, nb_classes)
    if _return_validation_data:
        assert X_val.shape == (nb_val_samples, nb_pixels_per_image)
        assert y_val.shape == (nb_val_samples, nb_classes)
    assert X_test.shape == (nb_test_samples, nb_pixels_per_image)
    assert y_test.shape == (nb_test_samples, nb_classes)

    # ====================================================================== #

    # step 6/6 : checking the datatype of the final formatted data

    check_dtype(X_train, utils.DEFAULT_DATATYPE)
    check_dtype(y_train, utils.DEFAULT_DATATYPE)
    if _return_validation_data:
        check_dtype(X_val, utils.DEFAULT_DATATYPE)
        check_dtype(y_val, utils.DEFAULT_DATATYPE)
    check_dtype(X_test, utils.DEFAULT_DATATYPE)
    check_dtype(y_test, utils.DEFAULT_DATATYPE)

    # ====================================================================== #

    t_end_data_formatting = perf_counter()
    duration_data_formatting = t_end_data_formatting - t_beginning_data_formatting

    # ====================================================================== #

    if verbose:
        print("\nShapes of the formatted MNIST data :")
        print(f"    - X_train : {X_train.shape}")
        print(f"    - y_train : {y_train.shape}")
        if _return_validation_data:
            print(f"    - X_val   : {X_val.shape}")
            print(f"    - y_val   : {y_val.shape}")
        print(f"    - X_test  : {X_test.shape}")
        print(f"    - y_test  : {y_test.shape}")

        print("\nTypes of the formatted MNIST data :")
        print(f"    - X_train : {get_dtype_of_array(X_train)}")
        print(f"    - y_train : {get_dtype_of_array(y_train)}")
        if _return_validation_data:
            print(f"    - X_val   : {get_dtype_of_array(X_val)}")
            print(f"    - y_val   : {get_dtype_of_array(y_val)}")
        print(f"    - X_test  : {get_dtype_of_array(X_test)}")
        print(f"    - y_test  : {get_dtype_of_array(y_test)}")

        precision_of_printed_info = 3
        print("\nRanges of the formatted MNIST data :")
        print(f"    - X_train : {get_range_of_array(X_train, precision=precision_of_printed_info)} (mean={X_train.mean():.{precision_of_printed_info}f}, std={X_train.std():.{precision_of_printed_info}f})")
        print(f"    - y_train : {get_range_of_array(y_train, precision=precision_of_printed_info)} (one-hot encoded)")
        if _return_validation_data:
            print(f"    - X_val   : {get_range_of_array(X_val,   precision=precision_of_printed_info)} (mean={X_val.mean():.{precision_of_printed_info}f}, std={X_val.std():.{precision_of_printed_info}f})")
            print(f"    - y_val   : {get_range_of_array(y_val,   precision=precision_of_printed_info)} (one-hot encoded)")
        print(f"    - X_test  : {get_range_of_array(X_test,  precision=precision_of_printed_info)} (mean={X_test.mean():.{precision_of_printed_info}f}, std={X_test.std():.{precision_of_printed_info}f})")
        print(f"    - y_test  : {get_range_of_array(y_test,  precision=precision_of_printed_info)} (one-hot encoded)")

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

        # displaying the class distributions of the final formatted data

        if _return_validation_data:
            dict_of_label_vectors = {
                "y_train" : categorical_to_vector(y_train, enable_checks=True),
                "y_val"   : categorical_to_vector(y_val,   enable_checks=True),
                "y_test"  : categorical_to_vector(y_test,  enable_checks=True)
            }
        else:
            dict_of_label_vectors = {
                "y_train" : categorical_to_vector(y_train, enable_checks=True),
                "y_test"  : categorical_to_vector(y_test,  enable_checks=True)
            }

        display_class_distributions(
            dict_of_label_vectors,
            selected_classes=selected_classes,
            dict_of_real_class_names=dict_of_real_class_names,
            precision=2
        )

    # ====================================================================== #

    if not(_return_validation_data):
        # checking that `X_val` and `y_val` have not been defined yet (within
        # this scope)
        assert ("X_val" not in locals()) and ("y_val" not in locals())

        # by design
        X_val, y_val = None, None

    print(f"\nThe raw MNIST dataset was successfully formatted. Done in {duration_data_formatting:.3f} seconds")

    return X_train, y_train, X_val, y_val, X_test, y_test


# import libraries
import numpy as np
import cv2

arrayCoords = []
def mouseLocationClick(event,x,y,flags,param):
    if event == cv2.EVENT_LBUTTONDOWN:
        print("click identified at: " +str([x,y]))
        arrayCoords.append([x,y])

imgRaw = cv2.imread("62,5.tif",0) # import the raw image here, currently set as "0,488.tif"

cv2.namedWindow('Raw Image') # make a named window
cv2.setMouseCallback('Raw Image', mouseLocationClick)
cv2.imshow('Raw Image',imgRaw)
keyPress = cv2.waitKey(0)
print(KeyPress)
if keyPress == 120:
    arrayBotRigCoords = arrayCoords.pop()
    arrayTopLefCoords = arrayCoords.pop()
    print("subplot coordinates: " + str(arrayBotRigCoords)+ " " + str(arrayTopLefCoords))
cv2.destroyAllWindows()

cropXCoords = sorted([arrayBotRigCoords[0],arrayTopLefCoords[0]])
cropYCoords = sorted([arrayBotRigCoords[1],arrayTopLefCoords[1]])
print(str(cropXCoords))
print(str(cropYCoords))
subImg = imgRaw[cropYCoords[0]:cropYCoords[1],cropXCoords[0]:cropXCoords[1]]
cv2.imshow('subcropped image',subImg)
cv2.waitKey(0)
cv2.destroyAllWindows()


#!/usr/bin/python
# coding=utf-8
from __future__ import division, print_function, unicode_literals
import imp
import os
import sys

BUILTINS = set(name for name in sys.builtin_module_names)

BUILTINS = BUILTINS.union(set(name.strip('_') for name in BUILTINS
                              if not name.startswith('__')))

def _postprocess_location(filename):
    if os.path.isdir(filename):
        filename = os.path.join(filename, '__init__.py')
    return filename

def locate_module(module_name, path=''):
    # step 1: Builtins
    if module_name in BUILTINS:
        return '__builtin__', False

    # step 2: local files and packages
    modparts = module_name.split('.')
    current_dir = os.path.abspath(path)
    found = True
    for modpart in modparts[:-1]:
        dirname = os.path.join(current_dir, modpart)
        if os.path.exists(dirname) and os.path.isdir(dirname) and\
                os.path.exists(os.path.join(dirname, '__init__.py')):
            current_dir = dirname
        else:
            found = False
            break

    if found:
        filename = os.path.abspath(os.path.join(current_dir, modparts[-1],
                                                '__init__.py'))
        if os.path.exists(filename):
            return _postprocess_location(filename), True

        filename = os.path.join(current_dir, modparts[-1] + '.py')
        if os.path.exists(filename):
            return _postprocess_location(filename), True

    # step 3: 3rd party modules
    try:
        f, filename, _ = imp.find_module(module_name)
        if f:
            f.close()
        #local = filename.startswith(path)
        return _postprocess_location(filename), False
    except ImportError:
        pass


# * codig: utf8 *
class TexWriter:
    def __init__(self):
        self.paperwidth = 10.0
        self.paperheight = 7.0

        self.lmargin = 0.3
        self.rmargin = 0.3
        self.bmargin = 0.3
        self.tmargin = 0.3

        self.lang = None

        self.titles = ['part', 'chapter', 'section', 'subsection', 'subsubsection', 'paragraph', 'subparagraph']
    def start(self):
        print("""
\\documentclass[oneside]{book}
\\usepackage[paperheight=%fcm,paperwidth=%fcm,tmargin=%fcm,rmargin=%fcm,bmargin=%fcm,lmargin=%fcm,heightrounded]{geometry}
\\usepackage[utf8]{inputenc}

%s
\\usepackage{indentfirst}
\\usepackage{misccorr}

\\usepackage{dejavu}

\\usepackage{graphicx}
\\DeclareGraphicsExtensions{.jpg,.png,.gif}

\\sloppy

\\usepackage{atbegshi}%% http://ctan.org/pkg/atbegshi
\\AtBeginDocument{\AtBeginShipoutNext{\AtBeginShipoutDiscard}}
        """ % (
            self.paperwidth,
            self.paperheight,
            self.tmargin,
            self.rmargin,
            self.bmargin,
            self.lmargin,
            ('\\usepackage[%s]{babel}' % self.lang) if self.lang != None else ''
        ))

    def meta(self, title, authors, date):
        print("\\author{%s}" % '\\and'.join(map(lambda author: self.encode(author), authors)))
        print("\\title{%s}" % self.encode(title))
        print("\\date{%s}" % self.encode(date))
        print("\\begin{document}")
        print("\\maketitle\n")

    def section(self, title, level = 0):
        print("\\%s{%s}" % (self.titles[level], self.encode(title.strip())))

    def text(self, content):
        print("%s\n" % self.encode(content));

    def end(self):
        print("\\end{document}")

    def encode(self, text):
        if (text is None):
            return ""
        # fix order
        escapeDict = [
            {'src': "\\", 'replacement': "\\textbackslash"},
            {'src': "{", 'replacement': "\\{"},
            {'src': "}", 'replacement': "\\}"},
            {'src': "&", 'replacement': "{\\&}"},
            {'src': "%", 'replacement': "\\%"},
            {'src': "$", 'replacement': "\\$"},
            {'src': "#", 'replacement': "\\#"},
            {'src': "_", 'replacement': "\\_"},
            {'src': "~", 'replacement': "\\texttt{\\~{}}"},
            {'src': "^", 'replacement': "\^{}"}
        ]
        for item in escapeDict:
            text = text.replace(item['src'], item['replacement'])
        return text


import collections

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

# Solution 1, simulation using hashmap
class Solution:
    def dfs(self, node, seen):
        if node is None:
            return
        seen[node.val] += 1

        self.dfs(node.left, seen)
        self.dfs(node.right, seen)

    def findMode(self, root: TreeNode) -> List[int]:
        if not root:
            return []
        seen = collections.Counter()

        self.dfs(root, seen)

        val = max(seen.values())
        res = [c for c in seen if seen[c] == val]

        return res

# Solution 2, recursion, returning two values trick
# O(1) memory solution if we ignore call stack space
class Solution:
    def dfs(self, node, prev_val, prev_cnt):
        if node is None:
            return prev_val, prev_cnt

        prev_val, prev_cnt = self.dfs(node.left, prev_val, prev_cnt)

        if prev_val == node.val:
            prev_cnt += 1
        else:
            prev_val, prev_cnt = node.val, 1

        if prev_cnt == self.res_cnt:
            self.res_vals.append(prev_val)
        elif prev_cnt > self.res_cnt:
            self.res_cnt = prev_cnt
            self.res_vals = [prev_val]

        return self.dfs(node.right, prev_val, prev_cnt)

    def findMode(self, root: TreeNode) -> List[int]:
        self.res_cnt = 0
        self.res_vals = []

        self.dfs(root, None, 0)

        return self.res_vals

Algorithms/0501 Find Mode in Binary Search Tree.py


import random
import os

def init():
    global started, dead, frames_while_dead_elapsed, ended, first_pipe, pipe_x, pipe_height, pipe_y, second_pipe_x, second_pipe_height, second_pipe_y, score, bird_angle, pipe_colour_green, unmovable
    started = False
    dead = False
    frames_while_dead_elapsed = 0
    ended = False
    first_pipe = True

    pipe_x = screen_width
    pipe_height = random.randint(25, 475)
    pipe_y = screen_height - pipe_height

    second_pipe_x = pipe_x + distance_between_pipes
    second_pipe_height = random.randint(25, 450)
    second_pipe_y = screen_height - second_pipe_height

    score = 0
    bird_angle = 0
    pipe_colour_green = random.choice([True, False])

    unmovable = False

def choose_pipe_colour():
    global pipe_colour_green
    pipe_colour_green = random.choice([True, False])

def set_bird_pictures():
    global bird_colour, wing_down, wing_mid, wing_up
    bird_colour = random.choice(["yellow", "blue", "red"])
    wing_up = os.path.join("data", "sprites", "birds", bird_colour, "up.png")
    wing_down = os.path.join("data", "sprites", "birds", bird_colour, "down.png")
    wing_mid = os.path.join("data", "sprites", "birds", bird_colour, "mid.png")

started = False
dead = False
frames_while_dead_elapsed = 0
ended = False
score = 0
unmovable = False

# -------------------------------------------------------------------------------------------------
fps_target = 62
screen_width, screen_height = 400, 600

# -------------------------------------------------------------------------------------------------
bird_colour = None
bird_x = screen_width * 0.4
bird_y = screen_height // 2
bird_change_y = 0
bird_height = 20
bird_width = 30
bird_angle = 0
bird_oscillation_angle = 0
bird_flap_state = 0
wing_up = None
wing_mid = None
wing_down = None

# -------------------------------------------------------------------------------------------------
distance_between_pipes = 300

# -------------------------------------------------------------------------------------------------
pipe_x = screen_width
pipe_height = random.randint(70, 450)
pipe_y = screen_height - pipe_height
pipe_width = 70
pipe_gap = 130

# -------------------------------------------------------------------------------------------------
second_pipe_x = pipe_x + distance_between_pipes
second_pipe_height = random.randint(70, 450)
second_pipe_y = screen_height - second_pipe_height

# -------------------------------------------------------------------------------------------------
pipe_velocity = -3.54
first_pipe = True
pipe_colour_green = random.choice([True, False])

if pipe_colour_green:
    top_pipe_path = os.path.join("data", "sprites", "pipes", "top_pipe.png")
    bottom_pipe_path = os.path.join("data", "sprites", "pipes", "bottom_pipe.png")
else:
    top_pipe_path = os.path.join("data", "sprites", "pipes", "top_pipe_red.png")
    bottom_pipe_path = os.path.join("data", "sprites", "pipes", "bottom_pipe_red.png")

# -------------------------------------------------------------------------------------------------
floor_width = 672
floor_x = 0
second_floor_x = floor_x + floor_width

# -------------------------------------------------------------------------------------------------
number0_path = os.path.join("data", "sprites", "numbers", "0.png")
number1_path = os.path.join("data", "sprites", "numbers", "1.png")
number2_path = os.path.join("data", "sprites", "numbers", "2.png")
number3_path = os.path.join("data", "sprites", "numbers", "3.png")
number4_path = os.path.join("data", "sprites", "numbers", "4.png")
number5_path = os.path.join("data", "sprites", "numbers", "5.png")
number6_path = os.path.join("data", "sprites", "numbers", "6.png")
number7_path = os.path.join("data", "sprites", "numbers", "7.png")
number8_path = os.path.join("data", "sprites", "numbers", "8.png")
number9_path = os.path.join("data", "sprites", "numbers", "9.png")


def maxDConnectedArea(G):
    '''
    计算有向图的最大连通区域
    :param G:example [[set(),set([1,2])],[set(0),set()],[set(0),set()]]
    '''
    n=len(G)
    def dfs(i):
        if G[i][1]:
            for j in G[i][1]:
                if j not in vis:
                    vis.add(j)
                    dfs(j)
    ans=0
    for i in range(n):
        vis = {i}
        dfs(i)
        ans=max(ans,len(vis))
    return ans


# Exercise 1 - Find Gain for Transparent Mode

import json as js
import gnpy.core.info as gn
import matplotlib.pyplot as plt
import numpy as np
import Mylib as ml
import gnpy.core.elements as gnel
import utilities as ut

# 2. 2. create a json file with the parameters of your Fiber. This json file has the
# following parameters:

#• ”uid”: a string used as unique identifier. You can use the name you
#want.
#• ”params”: a structure containing the following parameters:
#– ”length”: 80 km (length of the fiber)
#– ”loss coef”: 0.2 dB/km (attenuation coefficient)
#– ”length units”: ”km” (this parameter automatically scales the
#length and the loss coefficient)
#– ”att in”: 0 dB (attenuation before the fiber)
#– ”con in”: 0.5 dB (loss of the connector at the input of the fiber)
#– ”con out”: 0.5 dB (loss of the connector at the output of the
#fiber)
#– ”type variety”: ”SSMF”
#– ”dispersion”: 1.67e-05 s/m2
#(dispersion coefficient)
#– ’gamma’: 0.0 W/m (non-linear coefficient)

params = {"length": 80, "loss_coef": 0.2, "length_units": "km",
           "att_in" : 0, "con_in": 0.5, "con_out":0.5,
          "type_variety" : "SSMF", "dispersion": 1.67e-5,
          "gamma": 0.0 }

data = {"uid" : "fiber_id", "params" : params}

# write json data into a file
with open("Fiber_Parameters.json", "w") as write_file:
    js.dump(data, write_file, indent=4)

# 3. Instantiate the fiber from the JSON file.

# read json data from a file
with open("Fiber_Parameters.json", "r") as read_file:
    data = js.load(read_file)

fiber = gnel.Fiber(**data)

print(fiber)

# 4. Instantiate the spectral information.

# read json data from a file
with open("eqp.json", "r") as read_file:  # Exercise 2: ../eqp2.json"
    data = js.load(read_file)

# From the json data, extract only the values needed for the spectral info (SI)
si_data = data['SI'][0]

# Compute the WDM for the interested Spectral info
si = gn.create_input_spectral_information(si_data['f_min'], si_data['f_max'], si_data['roll_off'], si_data['baud_rate'], 10**(si_data['power_dbm']/10)/1000, si_data['spacing'])

# 5. Propagate the WDM comb through the Fiber.

propagated_si = fiber(si)

# 6. Calculate the loss for the channel 45

channels = propagated_si.carriers

initial_power_dBm = -1

output_fiber_power_W = channels[45].power.signal

output_fiber_power_dBm = 10*np.log10(output_fiber_power_W/1e-3)

print(output_fiber_power_dBm)

# output power should be Ptx - L*att - con_in - con_out = -1 -80*0.2 -0.5 -0.5

loss = initial_power_dBm - output_fiber_power_dBm

print("The loss is: ", loss)

# 7. Instantiate the amplifier, as in Exercise 1.4 Lesson 4, in such a way it
# recovers the loss.

params = {"gain_target": loss, "tilt_target": 0, "out_voa": 0}  # Exercise 2: "gain_target": 17

data = {"uid" : "amp_id", "type" : "Edfa", "type_variety" : "simple edfa", "operational" : params}

# write json data into a file
with open("Amp_Parameters.json", "w") as write_file:
    js.dump(data, write_file, indent=4)

Edfa_par = ut.get_edfa_parameters("Amp_Parameters.json","/Users/Tuna/PycharmProjects/OpenOptical/Lesson_6/eqp.json") # Exercise 2: ../eqp2.json"

EDFA = gnel.Edfa(**Edfa_par)

# propagate the WDM comb through the EDFA.

ml.plot_signal(propagated_si)

amplified_si = EDFA(propagated_si)

# 8. Plot the input and output signal power for each channel and verified the
# are the same.

ml.plot_signal(amplified_si)


#!/usr/bin/env python

import matplotlib.pyplot as plt
import numpy as np
from lightcurve_fitting.speccal import readspec
from corner import corner
from astropy.table import hstack, Table
from astropy.time import Time
from astropy.constants import c
import argparse
from matplotlib.gridspec import GridSpec
from fitting import MCgauss, gaussian

def closestpt(x, y, refx, refy, ax=None):
    """
    Returns the index of the point in a plotted Line2D that is closest to a mouse click on the plot axes

    Parameters
    ----------
    x : float
        x co-ordinate (in pixels) of the mouse click

    y : float
        y co-ordinate (in pixels) of the mouse click

    refx : np.ndarray
        Numpy array of wavelength values

    refy : np.ndarray
        Numpy array of flux values

    ax : matplotlib.axes.Axes, optional
        Axes of the Line2D plot. Default is current axes.

    Returns
    -------
    i : int
        Index into (refx, refy) that is closest to (x, y) in pixel space
    """
    if ax is None:
        ax = plt.gca()
    pixx, pixy = ax.transData.transform(list(zip(refx, refy))).T
    dist = ((x - pixx) ** 2 + (y - pixy) ** 2) ** 0.5
    i = np.argmin(dist)
    return i

class LinePlot:
    """
    Plot of a spectrum where the fitting region can be chosen interactively, with MCMC trace plots shown below

    Parameters
    ----------
    line : matplotlib.lines.Line2D
        Matplotlib Line2D object containing the spectrum

    refx : np.ndarray
        Numpy array of wavelength values

    refy : np.ndarray
        Numpy array of flux values

    ax2 : matplotlib.axes._subplots.AxesSubplot
        Subplots containing the MCMC traces during burn-in

    ax3 : matplotlib.axes._subplots.AxesSubplot
        Subplots containing the MCMC traces after burn-in

    profile : str
        Type of profile (input by the user)

    linewl : float
        Wavelength of line

    otherwl : float
        Wavelength of second line

    corner : bool, optional
        Makes corner plots if true. Default is False

    Attributes
    ----------
    line : matplotlib.lines.Line2D
        Matplotlib Line2D object containing the spectrum

    refx : np.ndarray
        Numpy array of wavelength values

    refy : np.ndarray
        Numpy array of flux values

    press : bool
        Indicates whether the mouse button is currently pressed

    i0 : int
        Index of (refx, refy) closest to where the mouse button was pressed

    i1 : int
        Index of (refx, refy) closest to where the mouse button was released

    ax2 : matplotlib.axes._subplots.AxesSubplot
        Subplots containing the MCMC traces during burn-in

    ax3 : matplotlib.axes._subplots.AxesSubplot
        Subplots containing the MCMC traces after burn-in

    fits : list
        Line2D objects plotted during the MCMC fitting. These are cleared after each mouse click.

    params : list
        Measurements from the MCMC fit: [mean(velocity), std(velocity), mean(EW), std(EW), mean(flux), std(flux)]

    profile : str
        Type of profile (input by the user)

    linewl : float
        Wavelength of line

    otherwl : float
        Wavelength of second line

    corner : bool
        Makes corner plots if true

    Methods
    -------
    connect()
        connect to all the events needed
    disconnect()
        disconnect all stored connection ids
    """

    def __init__(self, line, refx, refy, ax2, ax3, profile, linewl, otherwl, corner=False):
        self.line = line
        self.refx = refx
        self.refy = refy
        self.press = False
        self.i0 = None
        self.i1 = None
        self.ax2 = ax2
        self.ax3 = ax3
        self.fits = []
        self.params = [np.nan, np.nan, np.nan, np.nan, np.nan, np.nan]
        self.profile = profile
        self.linewl = linewl
        self.otherwl = otherwl
        self.corner = corner

    def connect(self):
        """
        Connect to all the events we need.

        """
        self.cidpress = self.line.figure.canvas.mpl_connect(
            'button_press_event', self._on_press)
        self.cidrelease = self.line.figure.canvas.mpl_connect(
            'button_release_event', self._on_release)
        self.cidmotion = self.line.figure.canvas.mpl_connect(
            'motion_notify_event', self._on_motion)

    def _on_press(self, event):
        """on button press we will see if the mouse is over us and store some data"""
        if event.inaxes != self.line.axes: return
        self.press = True
        self.i0 = closestpt(event.x, event.y, self.refx, self.refy, self.line.axes)
        self.line.set_xdata([self.refx[self.i0]] * 2)
        self.line.set_ydata([self.refy[self.i0]] * 2)
        for ax in self.ax2[:-1] + self.ax3[:-1]:
            ax.cla()
            plt.setp(ax.get_xticklabels(), visible=False)
        self.ax2[-1].cla()
        self.ax3[-1].cla()
        for l in self.fits:
            self.line.axes.lines.remove(l)
        self.fits = []
        self.line.figure.canvas.draw()

    def _on_motion(self, event):
        """on motion we will move the line if the mouse is over us"""
        if not self.press: return
        x0 = self.line.get_xdata()
        y0 = self.line.get_ydata()
        x0[1] = event.xdata
        y0[1] = event.ydata
        self.line.set_xdata(x0)
        self.line.set_ydata(y0)
        self.line.figure.canvas.draw()

    def _on_release(self, event):
        """on release we reset the press data"""
        self.press = False
        self.i1 = closestpt(event.x, event.y, self.refx, self.refy, self.line.axes)
        x0 = self.line.get_xdata()
        y0 = self.line.get_ydata()
        x0[1] = self.refx[self.i1]
        y0[1] = self.refy[self.i1]
        self.line.set_xdata(x0)
        self.line.set_ydata(y0)
        self.line.figure.canvas.draw()

        i0 = min(self.i0, self.i1)
        i1 = max(self.i0, self.i1)
        x = self.refx[i0:i1]
        y = self.refy[i0:i1]
        sampler, pos, model = MCgauss(x, y, profile=self.profile, linewl=self.linewl, otherwl=self.otherwl)
        for i in range(len(self.ax2)):
            self.ax2[i].plot(sampler.chain[:, :, i].T, 'k', alpha=0.2)
        sampler.reset()
        sampler.run_mcmc(pos, 1000)
        for i in range(len(self.ax3)):
            self.ax3[i].plot(sampler.chain[:, :, i].T, 'k', alpha=0.2)

        parameters = {param_name: chain for param_name, chain in zip(model.param_names, sampler.flatchain.T)}
        width = parameters['stddev']
        center = parameters['mean']
        continuum = parameters['intercept']
        slope = parameters.get('slope', 0.)
        amplitude_1 = parameters.get('amplitude_1', 0.)
        center_1 = parameters.get('mean_1', 0.)

        y1 = y[:, np.newaxis]

        # subtract off the second emission line
        if 'two' in self.profile:
            y1 = y1 - gaussian(x[:, np.newaxis], amplitude_1, width, center_1, 0)
            self.fits += self.line.axes.plot(x, y1.mean(axis=-1), color='b')

        # subtract off the flat or linear continuum
        y1 = y1 - continuum - slope * (x[:, np.newaxis] - center)

        if 'pcyg' in self.profile:
            velocity = c * (1. - center_1 / center)  # absorption minimum
        else:
            velocity = 2. * np.sqrt(2. * np.log(2.)) * c * width / center  # FWHM of the emission line
        equivwidth = np.trapz(y1 / continuum, x, axis=0)
        flux = np.trapz(y1, x, axis=0)

        v = np.mean(velocity)
        dv = np.std(velocity)
        ew = np.mean(equivwidth)
        dew = np.std(equivwidth)
        f = np.mean(flux)
        df = np.std(flux)
        self.params = [v, dv, ew, dew, f, df]

        print('FWHM velocity = {:.3f} +/- {:.3f}'.format(v, dv))
        print('equivalent width = {:.3f} +/- {:.3f}'.format(ew, dew))
        print('flux = {:.3e} +/- {:.3e}'.format(f, df))

        ps = [sampler.flatchain[i] for i in np.random.choice(sampler.flatchain.shape[0], 100)]
        for p in ps:
            yfit = model(x, *p)
            l = self.line.axes.plot(x, yfit, color='k', alpha=0.05)
            self.fits.append(l[0])
            if 'two' in self.profile:
                if len(p) == 6:
                    A, s, rA, x0, y0, dx = p
                    m = 0
                elif len(p) == 7:
                    A, s, rA, m, x0, y0, dx = p
                yfitA = gaussian(x, A, s, x0, y0) - m * (x - x0)
                lA = self.line.axes.plot(x, yfitA, color='g', alpha=0.05)
                self.fits.append(lA[0])
                yfitB = gaussian(x, A * rA, s, x0 + dx, y0) - m * (x - x0)
                lB = self.line.axes.plot(x, yfitB, color='r', alpha=0.05)
                self.fits.append(lB[0])

        self.line.figure.canvas.draw()

        if self.corner:
            corner(sampler.flatchain, labels=['A', 's', 'm', 'x0', 'y0', 'dx'])

    def disconnect(self):
        """disconnect all the stored connection ids"""
        self.line.figure.canvas.mpl_disconnect(self.cidpress)
        self.line.figure.canvas.mpl_disconnect(self.cidrelease)
        self.line.figure.canvas.mpl_disconnect(self.cidmotion)

def plot_results(t, ycol='flux'):
    """
    Plot the results of the MCMC fit from the table

    Parameters
    ----------
    t : astropy.table.Table
        The astropy table with the results
    ycol : str, optional
        The column to plot on the vertical axis. Default: 'flux'.
    """

    plt.errorbar(t['phase'], t[ycol], t['d'+ycol], fmt='o')
    plt.xlabel('Phase')
    plt.ylabel(ycol.capitalize())
    plt.show()

def measure_specvels(spectra, profile, linewl, l2=0., viewwidth=20., corner=False):
    """
    Fits a line profile to the spectrum and measures velocity, equivalent width, and flux

    Parameters
    ----------
    spectra : astropy.table.Table
        Astropy table containing the spectra to be fitted

    profile : str
        Type of profile (input by the user)

    linewl : float
        Wavelength of line (units must match the input spectrum)

    l2 : float, optional
        Wavelength of second line, if fitting two lines (units must match the input spectrum).

    viewwidth : float, optional
        Width of the spectrum plot in wavelength (units must match the input spectrum). Default is 20.

    corner : bool, optional
        Makes corner plots if true. Default is False

    Returns
    -------
    tout : astropy.table.Table
        The astropy table with the results
    """

    tparams = Table(names=('v', 'dv', 'EW', 'dEW', 'flux', 'dflux'))
    if profile == 'two':
        nparams = 6
        figsize = (5, 15)
    elif profile == 'pcyg':
        nparams = 6
        figsize = (5, 15)
    elif profile == 'emis':
        nparams = 4
        figsize = (5, 10)
    elif profile == 'bc' or profile == 'rc':
        nparams = 5
        figsize = (5, 12.5)
    elif profile == 'pcygbc':
        nparams = 7
        figsize = (5, 15)
    elif profile == 'twobc':
        nparams = 7
        figsize = (5, 15)
    else:
        raise ValueError('unrecognized profile type')
    hr = [4] + [1] * nparams
    gs = GridSpec(nparams + 1, 2, height_ratios=hr)
    for spec in spectra:
        good = ~np.isnan(spec['flux'])
        wl = spec['wl'][good]
        flux = spec['flux'][good]
        yshown = flux[(wl > linewl - viewwidth) & (wl < linewl + viewwidth)]
        if len(yshown) == 0:
            print('out of range for', spec['filename'], '(phase = {:+.1f})'.format(spec['phase']))
            tparams.add_row([np.nan, np.nan, np.nan, np.nan, np.nan, np.nan])
            continue
        f = plt.figure(figsize=figsize)
        ax1 = plt.subplot(gs[0:2])
        ax1.step(wl, flux)
        ax1.axvline(linewl, color='k', linestyle='-.', zorder=0)
        ax1.axis([linewl - viewwidth, linewl + viewwidth, np.min(yshown), np.max(yshown)])
        ax1.xaxis.tick_top()
        ax1.xaxis.set_ticks_position('both')

        l = plt.Line2D([], [])
        ax1.add_artist(l)

        ax2 = [plt.subplot(gs[2])]
        ax3 = [plt.subplot(gs[3])]
        for row in range(2, nparams + 1):
            ax2.append(plt.subplot(gs[2 * row], sharex=ax2[0]))
            ax3.append(plt.subplot(gs[2 * row + 1], sharex=ax3[0]))
        for ax in ax3:
            ax.yaxis.tick_right()
            ax.yaxis.set_ticks_position('both')
        for ax in ax2[:-1] + ax3[:-1]:
            plt.setp(ax.get_xticklabels(), visible=False)
        plt.subplots_adjust(left=0.11, right=0.89, top=0.97, bottom=0.03, hspace=0, wspace=0)

        dr = LinePlot(l, wl, flux, ax2, ax3, profile, linewl, l2, corner)
        dr.connect()
        plt.show()

        # cont = input('Press enter to accept these values or r to reject them.\n')
        # plt.clf()
        # if dr.corner:
        #     plt.figure(1)
        #     plt.clf()
        # if cont == 'r':
        #     print('fit rejected')
        #     tparams.add_row([np.nan, np.nan, np.nan, np.nan, np.nan, np.nan])
        # else:
        tparams.add_row(dr.params)

    for col in tparams.colnames:
        tparams[col].format = '%.3e'
    tout = hstack([spectra[['filename', 'date', 'telescope', 'instrument', 'phase']], tparams])
    tout[['filename', 'v', 'EW', 'flux']].pprint(max_width=-1, max_lines=-1)
    plot_results(tout)
    return tout

def read_spectra(filenames, redshift, refmjd=0.):
    """
    Read spectra from file and convert to rest wavelength

    Parameters
    ----------
    filenames : list
        List of filenames of the spectra to be read

    redshift : float
        Redshift to correct the spectra

    refmjd : float, optional
        Reference MJD to calculate phase. Default: 0.

    Returns
    -------
    spectra : astropy.table.Table
        Astropy table containing the spectra to be fitted
    """

    wls = []
    fluxes = []
    dates = []
    telescopes = []
    instruments = []
    for fn in filenames:
        wl, flux, date, tel, inst = readspec(fn)
        wl /= (1 + redshift)  # correct to rest wavelength
        if wl[0] > wl[-1]:  # if the spectrum is stored backwards, flip it now
            wl = wl[::-1]
            flux = flux[::-1]
        wls.append(wl)
        fluxes.append(flux)
        dates.append(Time.now() if date is None else date)
        telescopes.append(tel)
        instruments.append(inst)
    if len(wls) and all([len(wl) == len(wls[0]) for wl in wls]):  # stupid astropy workaround
        filenames.append('fake')
        wls.append(np.array([]))
        fluxes.append(np.array([]))
        dates.append(Time.now())
        telescopes.append('fake')
        instruments.append('fake')
        spectra = Table([filenames, dates, telescopes, instruments, wls, fluxes],
                        names=['filename', 'date', 'telescope', 'instrument', 'wl', 'flux'],
                        dtype=[str, object, str, str, np.ndarray, np.ndarray])
        spectra = spectra[:-1]
    else:
        spectra = Table([filenames, dates, telescopes, instruments, wls, fluxes],
                        names=['filename', 'date', 'telescope', 'instrument', 'wl', 'flux'],
                        dtype=[str, object, str, str, np.ndarray, np.ndarray])
    spectra['phase'] = (spectra['date'].mjd - refmjd) / (1. + redshift)
    return spectra

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Measure line velocities and equivalent widths in spectra.')
    parser.add_argument('filenames', nargs='+', help='filenames of the spectra')
    parser.add_argument('-l', '--wavelength', type=float, default=6563., help="wavelength of line")
    parser.add_argument('--l2', type=float, default=0, help="wavelength of second line")
    parser.add_argument('-z', '--redshift', type=float, default=0., help="redshift to correct the spectra")
    parser.add_argument('-t', '--refmjd', type=float, default=0., help="reference MJD to calculate phase")
    parser.add_argument('-w', '--viewwidth', type=float, default=200., help="estimated width of the line")
    parser.add_argument('-p', '--profile', default='emis', choices=['two', 'pcyg', 'emis', 'bc', 'rc', 'pcygbc', 'twobc'],
                        help="emission line (emis) or P Cygni line (pcyg)?")
    parser.add_argument('-c', '--corner', action='store_true', help="show corner plot")
    args = parser.parse_args()

    spectra = read_spectra(args.filenames, args.redshift, refmjd=args.refmjd)
    tout = measure_specvels(spectra, args.profile, args.wavelength, args.l2, args.viewwidth, args.corner)
    savefile = input('Save results as... ')
    if savefile:
        tout.write(savefile, format='ascii.fixed_width_two_line')


colors = open('_colors.scss', 'r').read().splitlines()

exportFile = open('_export.module.scss', 'w')

template = """@import "./_colors";

:export {{
{}
}}"""

colorsFormatted = []

for color in colors:
	if color:
		color = color.split(':')[0]
		cssname = color
		words = color[1:].split('-')
		for i in range(1, len(words)):
			words[i] = words[i].title()

		jsname = "".join(words)
		colorsFormatted.append(f"\t{jsname}: {cssname};")

exportFile.write(template.format("\n".join(colorsFormatted)))


"""
The probability package includes probability-based ATE approaches.
In addition, the package comes with functions that are likely to be useful to probability-based algorithms.
This includes caching functionality, which collects all documents across all corpora that mention a particular token.
"""

import json
import os
import sys

path = os.path.join(os.path.dirname(__file__), '..', '..', '..')
if path not in sys.path:
    sys.path.append(path)

import ate
from objects.exportable import Exportable
from summarization.timeline import Timeline

def p(corpora, focus=None, cache=None):
    """
    Calculate the probability of tokens appearing in the corpus.
    The probability is computed in terms of all tokens.

    Apart from calculating the probability of single tokens, the joint probability can be calculated by providing tuples.

    .. note::

        The joint probability is the minimum count of any token in the joint set in each document.

    :param corpora: A corpus, or corpora, of documents.
                    If a string is given, it is assumed to be one corpus.
                    If a list is given, it is assumed to be a list of corpora.

                    .. note::

                        It is assumed that the corpora were extracted using the tokenizer tool.
                        Therefore each line should be a JSON string representing a document.
                        Each document should have a `tokens` attribute.
    :type corpora: str or list of str
    :param focus: The tokens for which to compute the probability.
                 If nothing is given, the probability is calculated for all tokens.
                 The tokens can be provided as:

                 - A single word,
                 - A list of tokens,
                 - A tuple, or
                 - A list of tuples.

                 A tuple can be used to compute joint probabilities.
    :type focus: None or str or list of str or tuple or list of tuple
    :param cache: A list of terms that are re-used often and which should be cached.
                  If an empty list is given, no cache is used.

                  .. note::

                      Cache should be used when there is a lot of repetition.
                      For example, `x` can be used as cache when `x` is small and `y` is large.
                      If the data is small, using cache can be detrimental.
    :type cache: list of str

    :return: A dictionary with tokens as keys and probabilities as values.
    :rtype: dict
    """

    """
    Convert the corpora and tokens into a list if they aren't already.
    The list of tokens is always made into a list, even if it's a list of one string or tuple.
    """
    corpora = [ corpora ] if type(corpora) is str else corpora
    focus = focus or [ ]
    focus = [ focus ] if type(focus) is tuple or type(focus) is str else focus
    focus = [ (itemset, ) if type(itemset) is str else itemset for itemset in focus ]
    cache = cache or [ ]
    cache = [ cache ] if type(cache) is str else cache

    """
    Create the initial counts for all tokens and joint probabilities.
    This avoids returning missing probabilities for tokens or joint tokens that never appear.
    """
    counts = { }
    for itemset in focus:
        counts[itemset if len(itemset) > 1 else itemset[0]] = 0

    """
    If cache is defined, generate a list of documents for each cached token.
    This reduces the number of documents to go over.
    """
    if cache:
        for token in cache:
            """
            If no itemset contains the token, skip the cache.
            """
            if not any( token in itemset for itemset in focus ):
                continue

            """
            Create the cache.
            """
            documents = cached(corpora, token)

            """
            Look for item sets that mention the cached token.
            """
            itemsets = [ tuple for tuple in focus if token in tuple ]
            for document in documents:
                for itemset in itemsets:
                    if not all( item in document['tokens'] for item in itemset ):
                        continue

                    min_count = min( document['tokens'].count(item) for item in itemset )
                    counts[itemset if len(itemset) > 1 else itemset[0]] += min_count

            focus = [ tuple for tuple in focus if token not in tuple ]

    """
    Count the total number of tokens encountered and a separate count for each token.
    """
    tokens = 0
    for corpus in corpora:
        with open(corpus, 'r') as f:
            for line in f:
                document = json.loads(line)

                """
                The number of tokens is always incremented by all the tokens.
                """
                tokens += len(document['tokens'])

                """
                If there is no specification for which tokens to compute probability, compute the prior probability for all tokens.
                """
                if not focus:
                    for token in document['tokens']:
                        counts[token] = counts.get(token, 0) + 1
                else:
                    """
                    Convert each item in the list of tokens for which to compute the probability into a tuple.
                    """
                    for itemset in focus:
                        if not all( item in document['tokens'] for item in itemset ):
                            continue

                        min_count = min( document['tokens'].count(item) for item in itemset )
                        counts[itemset if len(itemset) > 1 else itemset[0]] += min_count

    """
    Compute the probability by dividing the count by the number of tokens.
    """
    if tokens:
        return { token: count / tokens for token, count in counts.items() }

    return { }

def joint_vocabulary(x, y):
    """
    Get the joint vocabulary by creating the cross-product from `x` and `y`.

    :param x: The tokens for which to compute the probability.
              These tokens are combined as a cross-product with all tokens in `y`.
              The tokens can be provided as:

              - A single word,
              - A list of tokens,
              - A tuple, or
              - A list of tuples.

              A tuple translates to joint probabilities.
    :type x: str or list of str or tuple or list of tuple
    :param y: The tokens for which to compute the probability.
              These tokens are combined as a cross-product with all tokens in `x`.
              The tokens can be provided as:

              - A single word,
              - A list of tokens,
              - A tuple, or
              - A list of tuples.

              A tuple translates to joint probabilities.
    :type y: str or list of str or tuple or list of tuple

    :return: The joint vocabulary, which is the cross-product of each item in `x` with each item in `y`.
    :rtype: list of tuple
    """

    vocabulary = [ ]

    """
    The list of tokens in `x` and `y` is always made into a list, even if it's a list of one string or tuple.
    """
    x, y = x or [ ], y or [ ]
    x = [ x ] if type(x) is tuple or type(x) is str else x
    y = [ y ] if type(y) is tuple or type(y) is str else y

    """
    Immediately return if either `x` or `y` are empty.
    """
    if not x:
        return [ tuple(item) if type(item) is not str else (item, ) for item in y ]

    if not y:
        return [ tuple(item) if type(item) is not str else (item, ) for item in x ]

    """
    Create the vocabulary.
    """
    for i in x:
        """
        Always convert the elements into a list, whether they are a string or a tuple.
        """
        i = list([ i ] if type(i) is str else i)
        for j in y:
            j = list([ j ] if type(j) is str else j)
            vocabulary.append(tuple(i + j))

    return vocabulary

def cached(corpora, token):
    """
    Compile all the documents in the given corpora that mention the token.
    These documents can be used as cache.
    In this way, the files do not have to be re-opened and documents without the token do not have to be iterated over.

    :param corpora: A corpus, or corpora, of documents.
                    If a string is given, it is assumed to be one corpus.
                    If a list is given, it is assumed to be a list of corpora.

                    It is assumed that the corpora are one of two types:

                    1. Corpora extracted using the :mod:`~tools.tokenizer` tool.
                       Therefore each line should be a JSON string representing a document.
                       Each document should have a ``tokens`` attribute.

                    2. Corpora generated using the :mod:`~tools.consume` tool.
                       Each such file is made up of a single :class:`~summarization.timeline.Timeline`.
                       Currently, only :class:`timelines <summarization.timeline.Timeline>` made up of :class:`~summarization.timeline.nodes.topical_cluster_node.TopicalClusterNode` instances are supported.
                       In this case, the function returns the number of nodes in the timeline.
    :type corpora: str or list of str
    :param token: The token to look for in the documents.
    :type token: str

    :return: A list of documents, each represented as a dictionary or a node, that contain the given token.
    :rtype: list of dict or list of :class:`~summarization.timeline.nodes.Node`
    """

    """
    Convert the corpora into a list if they aren't already.
    """
    corpora = [ corpora ] if type(corpora) is str else corpora

    documents = [ ]
    for corpus in corpora:
        with open(corpus, 'r') as f:
            if ate.datatype(corpus) is Timeline:
                """
                The check for timeline documents only has one step.
                It goes through each document and checks whether the token appears in the document.
                This is because the documents are already loaded.
                """
                timeline = Exportable.decode(json.loads(f.readline()))['timeline']
                for document in [ document for node in timeline.nodes for document in node.get_all_documents() ]:
                    if token in document.dimensions:
                        documents.append(document)
            else:
                for line in f:
                    """
                    The check for normal documents is in two steps:

                        1. Check whether the token appears in the line string; and
                        2. If it appears in the line string, decode the line and double-check that the token is in the document tokens.

                    This saves a lot of time from needlessly decoding lines, most of which do not contain the token.
                    """
                    if token in line:
                        document = json.loads(line)
                        if token in document['tokens']:
                            documents.append(document)

    return documents


"""https://www.hackerrank.com/challenges/special-palindrome-again/problem?h_l=interview&playlist_slugs%5B%5D%5B%5D=interview-preparation-kit&playlist_slugs%5B%5D%5B%5D=strings

A string is said to be a special string if either of two conditions is met:

All of the characters are the same, e.g. aaa.
All characters except the middle one are the same, e.g. aadaa.
A special substring is any substring of a string which meets one of those criteria. Given a string, determine how many special substrings can be formed from it.
"""

def is_special(s):
    """checks if the substring is "special"
    """
    # special case:
    if len(s) == 1: return True
    # general case
    match = s[0]
    for i in range(len(s)//2):
        if s[i] != match or s[i] != s[-(1+i)]: return False
    return True

def substrCount(n, s):
    """counts how many substrings are "special"

    somehow not fast enought... maybe because of the function call
    """
    n_specials = 0
    for i in range(len(s)):
        for j in range(i, len(s)):
            n_specials += 1 if is_special(s[i:j+1]) else 0
    return n_specials

def substrCount(n, s):
    res = 0
    count_sequence = 0
    prev = ''
    for i,v in enumerate(s):
        # first increase counter for all seperate characters
        count_sequence += 1
        if i and (prev != v):
            # if this is not the first char in the string
            # and it is not same as previous char,
            # we should check for sequence x.x, xx.xx, xxx.xxx etc
            # and we know it cant be longer on the right side than
            # the sequence we already found on the left side.
            j = 1
            while ((i-j) >= 0) and ((i+j) < len(s)) and j <= count_sequence:
                # make sure the chars to the right and left are equal
                # to the char in the previous found squence
                if s[i-j] == prev == s[i+j]:
                    # if so increase total score and step one step further out
                    res += 1
                    j += 1
                else:
                    # no need to loop any further if this loop did
                    # not find an x.x  pattern
                    break
            #if the current char is different from previous, reset counter to 1
            count_sequence = 1
        res += count_sequence
        prev = v
    return res

# test
# s = 'asasd'
s = 'abcbaba'
n = 5
print(f'testing with: {s}')
print(f'ans: {substrCount(n, s)}')


# -*- coding: utf-8 -*-
"""
Created on Sun Apr 26 13:47:17 2020

@author: matis
"""

import networkx as nx

G=nx.complete_graph(6)
print (nx.coloring.greedy_color(G)) #greedy approach algorithm
#Los numeros del diccionario son los vértices
#Los números dentro de las llaves (vertices) del diccionario, son el color asignado
K=nx.gnm_random_graph(10,12)
print(nx.coloring.greedy_color(K))


"""
.. index:: Extensions; table_styling

==============================
PSphinxTheme.ext.table_styling
==============================

Overview
========
**adds directives for styling tables**

This Sphinx extension replaces the default ``.. table::`` directive with a custom one, that supports a number of extra
options for controlling table layout on a per-column basis.

.. index:: Tables; table_styling Table examples, Usage; extension: table_styling usage example

For example, the following snippet specifies relative widths for the three columns, changes the text alignment for each
column, disables text-wrapping for the third column, and inserts dividers between the columns::

   .. table:: Optional Caption
      :widths: 3 2 1
      :column-alignment: left center right
      :column-wrapping: true true false
      :column-dividers: none single double single

      =========== =========== ===========
      Width 50%   Width 33%   Width 16%
      =========== =========== ===========
      Line 1      This text   This text
                  should wrap will always
                  onto        be one line.
                  multiple
                  lines.
      Line 2      Centered.   Right-Aligned.
      Line 3      Centered    Right-Aligned
                  Again.      Again.
      =========== =========== ===========

This will then result in the following table:

.. table:: Optional Caption
   :widths: 3 2 1
   :column-alignment: left center right
   :column-wrapping: yes yes no
   :column-dividers: none single double single

   =========== =========== ===========
   Width 50%   Width 33%   Width 16%
   =========== =========== ===========
   Line 1      This text   This text
               should wrap will always
               onto        be one line.
               multiple
               lines.
   Line 2      Centered.   Right-Aligned.
   Line 3      Centered    Right-Aligned
               Again.      Again.
   =========== =========== ===========

Directive Options
=================
The enhanced ``.. table::`` directive supports the following options:

``:widths:``
   Sets proportional column widths

   This should be a space-separated list of positive integers, one for every column. The columns widths will be allocated
   proportionally (e.g. ``1 2 3`` for a 3-column table means the columns will use 16%, 33%, and 50% of the total width,
   respectively).

``:column-alignment:``
   Sets per-column text alignment

   This should be a comma/space-separated list of the following strings: ``left``, ``right``, ``center``, ``justify``.
   These are interpreted one per column. Extra values are ignored; if not enough values are provided, the remaining columns
   will default to ``left``.

   Alternately this can be a single word containing just the first letters: e.g. ``lrcj`` would be interpreted the same as
   ``left right center justify``.

``:column-wrapping:``
   Sets per-column text wrapping

   This should be a comma/space-separated list of either ``yes``/``true`` or ``no``/``false``. If ``yes`` (the default),
   words will wrap around if there is not enough horizontal space. If ``no``, whitespace-wrapping will be disabled for that
   column. Extra values are ignored; if not enough values are provided, the remaining columns will default to ``true``.

   Alternately this can be a single word containing just the first letters: e.g. ``ttf`` is the same as ``true true false``.

``:column-dividers:``
   Add dividers between columns

   This lets you specify custom vertical dividers between columns (ala what ``.. tabularcolumns::`` allows under latex).

   This should be a comma/space-separated list of ``none``, ``single``, or ``double``; based on the type of divider you want.
   There should be one of these for the left side of the table, for between each column, and for the right side of the table
   (e.g. a 4 entries for a 3-column table). Extra values are ignored; if not enough values are provided, the remaining
   columns will default to ``none``.

   Alternately this can be a single word containing just the number: e.g. ``0121`` is the same as
   ``none single double single``.

``:column-classes:``
   Add per-column css classes.

   This lets you specify css classes that will be assigned to each column, much like ``.. rst-class::``. This should either
   be a space-separated list containing one class per column, or a comma-separated list containing multiple classes for each
   column, separated by spaces. Extra values are ignored; if there are not enough values, or there are blank entries, those
   columns will not be assigned any custom classes.

   For example, ``a b, c , , d`` would assign the classes ``a b`` to column 1, ``c`` to column 2, no custom classes for
   column 3, and ``d`` to column 4.

``:header-columns:``
   Specify number of "stub" columns

   Should be a non-negative integer specifying the number of columns (starting with the left side) that should be treated as
   "stub" or "header" columns, and should be styled accordingly.

Internals
=========
By default, all HTML tables styled by this extension will have the css class ``"styled-table"`` added, to help with themeing
support. Use this option to override with your own theme's preferred css class.

.. note::

   This extension gets the job done by adding custom css classes to every cell in the generated html table. It then inserts a
   custom css file which provides styling implementing relevant parts of the above directives. While this extension is
   primarily tested with :mod:`!PSphinxTheme`, it should work with most Sphinx themes, any conflicts that occur are probably
   bugs.

.. todo:: allow ``:widths:`` to support ``em``, ``in``, ``%``

"""
from itertools import zip_longest

from docutils import nodes
from docutils.parsers.rst import directives
from docutils.parsers.rst.directives.tables import RSTTable

__all__ = ['setup', ]

# ===========================================================================================================================
# field option parsers
# ===========================================================================================================================
def _parse_argument_map(argument, argmap, param):
   """ Helper: _parse_argument_map
   """
   if ',' in argument:
      args = argument.split(',')
   else:
      args = argument.split()
   if len(args) == 1 and all(c in argmap for c in args[0]):
      args = args[0]

   def norm(arg):
      try:
         return argmap[arg]
      except KeyError:
         raise ValueError('invalid %s: %r' % (param, arg))

   return [norm(arg) for arg in args]

_alignment_map = dict(
   l='left',
   r='right',
   c='center',
   j='justify',
   left='left',
   right='right',
   center='center',
   centered='center',  # compat alias
   justify='justify',
   justified='justify',  # compat alias
)

def alignment_list(argument):
   """ convert into list of alignment options.
   raise ``ValueError`` if no args found, or invalid strings.
   """
   return _parse_argument_map(argument, _alignment_map, 'alignment')

_bool_map = {
   'true': True,
   't': True,
   'yes': True,
   'y': True,
   'false': False,
   'f': False,
   'no': False,
   'n': False,
}

def bool_list(argument):
   """ convert to list of true/false values
   """
   return _parse_argument_map(argument, _bool_map, 'boolean value')

def class_option_list(argument):
   """ convert to list of list of classes
   """
   if ',' in argument:
      args = argument.split(',')
   else:
      args = argument.split()
   return [directives.class_option(arg) for arg in args]

_divider_map = {
   '0': 'no',
   '1': 'single',
   '2': 'double',
   'none': 'no',
   'single': 'single',
   'double': 'double',
}

def divider_list(argument):
   return _parse_argument_map(argument, _divider_map, 'divider style')

# ===========================================================================================================================
# replacement for table directive
# ===========================================================================================================================
class ExtendedRSTTable(RSTTable):
   option_spec = RSTTable.option_spec.copy()
   option_spec.update({
      # class, name already present
      # 'header-rows': directives.nonnegative_int,
      'header-columns': directives.nonnegative_int,
      # TODO: column-widths: support limited set of units (em/in/%)
      # expressible under both css & latex
      'widths': directives.positive_int_list,
      'column-alignment': alignment_list,
      'column-wrapping': bool_list,
      'column-classes': class_option_list,
      'column-dividers': divider_list,
   })

   def run(self):
      """ run
      """
      result = RSTTable.run(self)
      if result and isinstance(result[0], nodes.table):
         self._update_table_classes(result[0])
      return result

   def _update_table_classes(self, table):
      """ _update_table_classes
      """
      assert isinstance(table, nodes.table)
      table['classes'].append('styled-table')
      header_cols = self.options.get('header-columns') or 0
      # header_rows = self.options.get('header-rows')
      widths = self.options.get('widths')
      dividers = self.options.get('column-dividers')
      if dividers is None:
         get_divider = None
      else:
         def get_divider(idx_):
            try:
               return dividers[idx_]
            except IndexError:
               return 'no'
      # noinspection PyPep8Naming
      EMPTY = ()
      opts = (
         self.options.get('column-alignment', EMPTY),
         self.options.get('column-wrapping', EMPTY),
         self.options.get('column-classes', EMPTY),
      )

      # noinspection PyShadowingNames
      def locate(cls):
         """ Helper: locate
         """
         for child_ in table.children:
            if isinstance(child_, cls):
               return child_
         return None

      tgroup = locate(nodes.tgroup)
      if not tgroup:
         return
      col = 0
      for child in tgroup:
         if isinstance(child, nodes.colspec):
            if widths and col < len(widths):
               child['colwidth'] = widths[col]
            if col < header_cols:
               child['stub'] = 1
            col += 1
            continue
         assert isinstance(child, (nodes.thead, nodes.tbody))
         for row in child.children:
            # add alignment and wrap classes to each entry
            # (would add to colspec, but html doesn't inherit much from colgroup)
            assert isinstance(row, nodes.row)
            for idx, (entry, align, wrap, clist) in enumerate(zip_longest(row, *opts)):
               if entry is None:
                  raise ValueError('not enough columns for field options')
               assert isinstance(entry, nodes.entry)
               classes = entry['classes']
               if align:
                  classes.append(align + '-align')
               if wrap is False:
                  classes.append('nowrap')
               if clist:
                  classes.extend(clist)
               if get_divider:
                  # noinspection PyCallingNonCallable
                  classes.append(get_divider(idx) + '-left-divider')
                  # noinspection PyCallingNonCallable
                  classes.append(get_divider(idx + 1) + '-right-divider')

# ===========================================================================================================================
# register extension
# ===========================================================================================================================
def setup(app):
   """ register extension
   """
   # replace existing table directive with custom one
   app.add_directive('table', ExtendedRSTTable)

# ===========================================================================================================================
# eof
# ===========================================================================================================================


import logging
import zstandard
import json
import os
import time

class ZstandardHandler():
    def __init__(self, t_info_storage):
        logging.basicConfig(format='%(asctime)s - %(message)s', level=logging.INFO)

        self.t_info_storage = t_info_storage

    # Function decompresses file_path and returns line by line (Generator)
    # Code used from: github.com/Watchful1/PushshiftDumps
    def read_file(self, inner_torrent_path, save_folder_path):
        file_save_path = save_folder_path + inner_torrent_path

        file_exists = False
        while file_exists == False:
            #print(f"Waiting for {file_save_path} to exist")
            if os.path.exists(file_save_path) == True:
                file_exists = True
                #print(f"The file ({file_save_path}) has been found")
            time.sleep(0.5)

        with open(file_save_path, 'rb') as file_handle:
            buffer = ''
            reader = zstandard.ZstdDecompressor(max_window_size=2**31).stream_reader(file_handle)
            total_waits = 0
            while True:
                torrent_info = self.t_info_storage.get_torrent_info(inner_torrent_path)
                # Progress is 0 to 1 representing how much of the file is downloaded, size is total size in bytes
                downloaded_total = torrent_info['progress'] *  torrent_info['size']

                # Change Multiplier to 2 if breaking?
                downloaded_total_needed = file_handle.tell() + ((2**27) * 2)

                if ( downloaded_total_needed > downloaded_total and torrent_info['progress'] != 1):
                    #logging.info(f"Waiting for more of {file_save_path} to download ({downloaded_total}/{downloaded_total_needed})")
                    total_waits += 1
                    time.sleep(2)
                else:
                    #logging.info(f"Reading a chunk of {file_save_path}")
                    total_waits = 0

                    chunk = self.read_and_decode(reader,
                                                chunk_size=2**27,
                                                max_window_size=(2**29) * 2)
                    # If nothing returned above then stop reading
                    # (I.e. End of file)
                    if not chunk:
                        break

                    # Adds last bit of previous chunk to start of this chunk
                    # Then splits by new line
                    lines = (buffer + chunk).split("\n")

                    # Returns all lines except last line
                    for line in lines[:-1]:
                        json_line = json.loads(line.strip())
                        # Returns JSON and amount of bits read through the file
                        yield json_line, file_handle.tell()

                    # Save the last line of this chunk as likely not complete so can add it to next chunk
                    buffer = lines[-1]

            reader.close()

    # Functions for uncompressing each chunk
    # Code used from: github.com/Watchful1/PushshiftDumps
    def read_and_decode(self, reader, chunk_size, max_window_size, previous_chunk=None, bytes_read=0):
        try:
            chunk = reader.read(chunk_size)
        except zstandard.ZstdError as e:
            logging.info("Error happened")
        bytes_read += chunk_size
        if previous_chunk is not None:
            chunk = previous_chunk + chunk
        try:
            return chunk.decode()
        except UnicodeDecodeError:
            if bytes_read > max_window_size:
                raise UnicodeError(f"Unable to decode frame after reading {bytes_read:,} bytes")
            logging.info(f"Decoding error with {bytes_read:,} bytes, reading another chunk")
            return self.read_and_decode(reader, chunk_size, max_window_size, chunk, bytes_read)


numberRoutines = int(input())
minQuickChanges = float('Inf')
dances = []
for i in range(0, numberRoutines):
    dance = input()
    dances.append(dance)
def countQuickChanges(routine):
    num = 0
    for i in range(1, len(routine)):
        chars = set()
        for char in routine[i-1]:
            chars.add(char)

        for char in routine[i]:
            if char in chars:
                num += 1
    return num

def chooseDancers(remaining_dances, routine, dance_index):
    global minQuickChanges
    # print(routine)
    if len(remaining_dances) == 0:
        quickChanges = countQuickChanges(routine)
        if quickChanges < minQuickChanges:
            minQuickChanges = quickChanges
        return True
    for i in range(len(remaining_dances)):
        routine[dance_index] = remaining_dances[i]
        # remaining_dances.remove(dance)
        dance = remaining_dances[i]
        del remaining_dances[i]
        chooseDancers(remaining_dances, routine, dance_index + 1)
        remaining_dances.insert(i, dance)

chooseDancers(dances, ['' for i in range(numberRoutines)], 0)
print(minQuickChanges)


from typing import Sequence

from absl import app
from absl import flags
from circuit_training.environment import plc_client

flags.DEFINE_string("netlist_file", None, "Path to the input netlist file.")

flags.mark_flags_as_required([
    "netlist_file",
])

FLAGS = flags.FLAGS

def main(argv: Sequence[str]) -> None:
  if len(argv) > 1:
    raise app.UsageError("Too many command-line arguments.")

  plc = plc_client.PlacementCost(netlist_file=FLAGS.netlist_file)

  print("get_cost:", plc.get_cost())
  print("get_congestion_cost:", plc.get_congestion_cost())
  print("get_density_cost:", plc.get_density_cost())

  hard_macro_indices = [
      m for m in plc.get_macro_indices() if not plc.is_node_soft_macro(m)
  ]
  print("hard_macro_indices:", hard_macro_indices)
  print("get_node_mask:", plc.get_node_mask(hard_macro_indices[0]))

if __name__ == "__main__":
  app.run(main)


from base_evaluator import BaseEvaluator

def is_most_freq_entity(entity_idx, entity_freqs):
    entity_freqs[entity_idx] -= 1
    most_freq_entity_idx = entity_freqs.index(max(entity_freqs))
    entity_freqs[entity_idx] += 1
    return entity_idx == most_freq_entity_idx

class MostFreqEntityEvaluator(BaseEvaluator):
    def __init__(self, logger=None, use_lemma=True, include_type=True,
                 include_all_pobj=True, ignore_first_mention=False,
                 filter_stop_events=True):
        super(MostFreqEntityEvaluator, self).__init__(
            logger=logger,
            use_lemma=use_lemma,
            include_type=include_type,
            include_all_pobj=include_all_pobj,
            ignore_first_mention=ignore_first_mention,
            filter_stop_events=filter_stop_events
        )
        self.model_name = 'most_freq_entity'

    def set_model(self, model):
        self.set_embedding_model(model)

    def evaluate_event_list(self, rich_event_list):
        for event_idx, rich_event in enumerate(rich_event_list):
            self.logger.debug('Processing event #{}'.format(event_idx))
            for arg_idx in rich_event.get_arg_idx_list(
                    include_all_pobj=self.include_all_pobj):
                if rich_event.has_neg(arg_idx):
                    rich_arg = rich_event.get_argument(arg_idx)
                    if not self.ignore_argument(rich_arg):
                        entity_freqs = [rich_entity.salience.num_mentions_total
                                        for rich_entity in
                                        rich_arg.rich_entity_list]
                        correct = is_most_freq_entity(
                            rich_arg.entity_idx, entity_freqs)
                        num_choices = len(entity_freqs)

                        kwargs = BaseEvaluator.get_arg_group_info(rich_arg)

                        self.eval_stats.add_eval_result(
                            correct,
                            num_choices,
                            **kwargs
                        )


import time
from smbus import SMBus
print(">>>>>tx_data module successfully imported")

addr = 0x8  # bus address
bus = SMBus(1)  # indicates /dev/ic2-1

def ConvertStringsToBytes(src):
    converted = []
    for b in src:
        converted.append(ord(b))
    return converted

def handshake():
    time.sleep(0.005)
    with SMBus(1) as bus:
        bus.write_byte_data(addr, 0, 2)
    try:
        data = bus.read_byte_data(addr, 0)
        return data
    except:
        print("remote i/o error")
        return -1

def drive(throttle, steering):
    time.sleep(0.005)
    with SMBus(1) as bus:
        bus.write_byte_data(addr, 0, 3)
    time.sleep(0.01)
    with SMBus(1) as bus:
        bus.write_byte_data(addr, 0, ConvertStringsToBytes(
            str(steering) + "/"+str(throttle)))
# check if the timing works out, adjust accordingly

def stop():
    time.sleep(0.005)
    with SMBus(1) as bus:
        bus.write_byte_data(addr, 0, 4)

def kill():
    time.sleep(0.005)
    with SMBus(1) as bus:
        bus.write_byte_data(addr, 0, 5)

def init():
    time.sleep(0.005)
    with SMBus(1) as bus:
        bus.write_byte_data(addr, 0, 6)