"""
Interface to NasBench101 for Hyperparameter Optimization and Neural Architecture Search

https://github.com/automl/nas_benchmarks

How to use this benchmark:
--------------------------

We recommend using the containerized version of this benchmark.
If you want to use this benchmark locally (without running it via the corresponding container),
you need to perform the following steps.

1. Download data
================

The data will be downloaded automatically.
Note: However, if you use the benchmark locally, you can specify also the data directory (path to the folder, where the
nasbench_full.tfrecord is) by hand.

In this case you can download the data with the following command.

```
wget https://storage.googleapis.com/nasbench/nasbench_full.tfrecord
```
Remark: it is important to select the full tf record and not the 'only_108' record to perform multi-fidelity
optimization.

2. Clone and install
====================
```
cd /path/to/HPOBench
pip install .[nasbench_101]

pip install git+https://github.com/google-research/nasbench.git@master
pip install git+https://github.com/automl/nas_benchmarks.git@master
```

Notes:
------
Benchmarks in NASBench101 only contain epochs 4, 12, 36 and 108.
Querying another epoch, e.g. 5, raises an assertion.

Changelog:
==========
0.0.4
* New container release due to a general change in the communication between container and HPOBench.
  Works with HPOBench >= v0.0.8

0.0.3:
* Standardize the structure of the meta information

0.0.2:
* The objective function takes as input now the parameter run_index. Allowed values are Tuple(0-2), 0, 1, 2, None.
  This value specifies which seeds are used. The user can specify a single index or a tuple with indices.
  If the user wants to use a randomly drawn run_index, they can simply set the value explicitly to None.
* Fix a bug in NASCifar10CBenchmark

0.0.1:
* First implementation

"""
import logging

from pathlib import Path
from typing import Union, Dict, Any, Tuple, List

import ConfigSpace as CS
import numpy as np
from tabular_benchmarks.nas_cifar10 import NASCifar10
from nasbench import api
from nasbench.api import OutOfDomainError
from nasbench.lib import graph_util

from hpobench import config_file
import hpobench.util.rng_helper as rng_helper
from hpobench.abstract_benchmark import AbstractBenchmark
from hpobench.util.data_manager import NASBench_101DataManager

__version__ = '0.0.4'
logger = logging.getLogger('NasBench101')

MAX_EDGES = 9
VERTICES = 7
DEFAULT_API_FILE = config_file.data_dir / "nasbench_101"

class NASCifar10BaseBenchmark(AbstractBenchmark):
    def __init__(self, benchmark: NASCifar10,
                 data_path: Union[Path, str, None] = None,
                 rng: Union[np.random.RandomState, int, None] = None, **kwargs):
        """
        Baseclass for the tabular benchmarks https://github.com/automl/nas_benchmarks/tree/master/tabular_benchmarks.
        Please install the benchmark first. Place the data under ``data_path``.

        Parameters
        ----------
        benchmark : NASCifar10
            Type of the benchmark to use. Don't call this class directly. Instantiate via subclasses (see below).
        data_path : str, Path, None
            Path to the folder, which contains the downloaded file nasbench_full.tfrecord.
        rng : np.random.RandomState, int, None
            Random seed for the benchmarks
        """

        super(NASCifar10BaseBenchmark, self).__init__(rng=rng)

        self.benchmark = benchmark
        self.data_path = data_path

    def _query_benchmark(self, config: Dict, run_index: int, budget: int = 108) -> Dict:
        raise NotImplementedError

    # pylint: disable=arguments-differ
    @AbstractBenchmark.check_parameters
    def objective_function(self, configuration: Union[CS.Configuration, Dict],
                           fidelity: Union[CS.Configuration, Dict, None] = None,
                           run_index: Union[int, Tuple, None] = (0, 1, 2),
                           rng: Union[np.random.RandomState, int, None] = None,
                           **kwargs) -> Dict:
        """
        Query the NAS-benchmark using a given configuration and a epoch (=budget).

        Parameters
        ----------
        configuration : Dict, CS.Configuration
        fidelity: Dict, None
            Fidelity parameters, check get_fidelity_space(). Uses default (max) value if None.
        run_index : int, Tuple, None
            The nas benchmark has for each configuration-budget-pair results from 3 different runs.
            - If multiple `run_id`s are given as Tuple, the benchmark returns the mean over the given runs.
            - By default (no parameter is specified) all runs are used. A specific run can be chosen by setting the
              `run_id` to a value from [0, 3]. While the performance is averaged across the `run_index`, the costs are
              the sum of the runtime per `run_index`.
            - When this value is explicitly set to `None`, the function will use a random seed.
        rng : np.random.RandomState, int, None
            Random seed to use in the benchmark.

            To prevent overfitting on a single seed, it is possible to pass a
            parameter ``rng`` as 'int' or 'np.random.RandomState' to this function.
            If this parameter is not given, the default random state is used.
        kwargs

        Returns
        -------
        Dict -
            function_value : validation error
            cost : runtime
            info : Dict
                fidelity : used fidelities in this evaluation
        """
        self.rng = rng_helper.get_rng(rng, self_rng=self.rng)

        if isinstance(run_index, int):
            assert 0 <= run_index <= 2, f'run_index must be in [0, 2], not {run_index}'
            run_index = (run_index, )
        elif isinstance(run_index, (Tuple, List)):
            assert 0 < len(run_index) <= 3, 'run_index must not be empty'
            assert min(run_index) >= 0 and max(run_index) <= 2, \
                f'all run_index values must be in [0, 2], but were {run_index}'
            if len(set(run_index)) != len(run_index):
                logger.debug('There are some values more than once in the run_index. We remove the redundant entries.')
                run_index = tuple(set(run_index))
        elif run_index is None:
            logger.debug('The run index is explicitly set to None! A random seed will be selected.')
            run_index = tuple(self.rng.choice((0, 1, 2), size=1))
        else:
            raise ValueError(f'run index must be one of Tuple or Int, but was {type(run_index)}')

        self.benchmark.reset_tracker()

        # Returns (valid_accuracy: 0, runtime: 0) if it is invalid, e.g. config not valid or
        # budget not in 4 12 36 108
        train_accuracies = []
        valid_accuracies = []
        test_accuracies = []
        training_times = []
        additional = {}

        for run_id in run_index:
            data = self._query_benchmark(config=configuration, budget=fidelity['budget'], run_index=run_id)

            train_accuracies.append(data['train_accuracy'])
            valid_accuracies.append(data['validation_accuracy'])
            test_accuracies.append(data['test_accuracy'])
            training_times.append(data['training_time'])

            # Since those information are the same for all run ids, just store one of them.
            additional = {'trainable_parameters': data['trainable_parameters'],
                          'module_operations': data['module_operations']}

        return {'function_value': float(1 - np.mean(valid_accuracies)),
                'cost': float(np.sum(training_times)),
                'info': {'fidelity': fidelity,
                         'train_accuracies': train_accuracies,
                         'valid_accuracies': valid_accuracies,
                         'test_accuracies': test_accuracies,
                         'training_times': training_times,
                         'data': additional
                         }
                }

    @AbstractBenchmark.check_parameters
    def objective_function_test(self, configuration: Union[Dict, CS.Configuration],
                                fidelity: Union[CS.Configuration, Dict, None] = None,
                                rng: Union[np.random.RandomState, int, None] = None,
                                **kwargs) -> Dict:
        """
        Validate a configuration on the maximum available budget.

        Parameters
        ----------
        configuration : Dict, CS.Configuration
        fidelity: Dict, None
            Fidelity parameters, check get_fidelity_space(). Uses default (max) value if None.
        rng : np.random.RandomState, int, None
            Random seed to use in the benchmark. To prevent overfitting on a single seed, it is
            possible to pass a parameter ``rng`` as 'int' or 'np.random.RandomState' to this
            function. If this parameter is not given, the default random state is used.
        kwargs

        Returns
        -------
        Dict -
            function_value : test error
            cost : runtime
            info : Dict
                fidelity : used fidelities in this evaluation
        """

        result = self.objective_function(configuration=configuration, fidelity=fidelity, run_index=(0, 1, 2), rng=rng)
        result['function_value'] = float(1 - np.mean(result['info']['test_accuracies']))

        return result

    @staticmethod
    def get_configuration_space(seed: Union[int, None] = None) -> CS.ConfigurationSpace:
        raise NotImplementedError

    @staticmethod
    def get_meta_information() -> Dict:
        """ Returns the meta information for the benchmark """
        return {'name': 'Tabular Benchmarks for Hyperparameter Optimization and Neural Architecture Search',
                'references': ['@article{klein2019tabular,'
                               'title   = {Tabular benchmarks for joint architecture and hyperparameter optimization},'
                               'author  = {Klein, Aaron and Hutter, Frank},'
                               'journal = {arXiv preprint arXiv:1905.04970},'
                               'year    = {2019}}',
                               'https://arxiv.org/abs/1905.04970',
                               ],
                'code': 'https://github.com/automl/nas_benchmarks',
                }

    @staticmethod
    def _get_configuration_space(benchmark: Any, seed: Union[int, None] = None) -> CS.ConfigurationSpace:
        """ Helper function to pass a seed to the configuration space """
        seed = seed if seed is not None else np.random.randint(1, 100000)
        cs = benchmark.get_configuration_space()
        cs.seed(seed)
        return cs

    @staticmethod
    def get_fidelity_space(seed: Union[int, None] = None) -> CS.ConfigurationSpace:
        """
        Creates a ConfigSpace.ConfigurationSpace containing all fidelity parameters for
        the NAS Benchmark 101.

        Parameters
        ----------
        seed : int, None
            Fixing the seed for the ConfigSpace.ConfigurationSpace

        Returns
        -------
        ConfigSpace.ConfigurationSpace
        """
        seed = seed if seed is not None else np.random.randint(1, 100000)
        fidel_space = CS.ConfigurationSpace(seed=seed)

        fidel_space.add_hyperparameters([
            CS.OrdinalHyperparameter('budget', sequence=[4, 12, 36, 108], default_value=108)
        ])

        return fidel_space

    @staticmethod
    def _try_download_api_file(save_to: Union[Path, str, None]):
        data_manager = NASBench_101DataManager(save_to)
        data_manager.download()
        return data_manager.save_dir

class NASCifar10ABenchmark(NASCifar10BaseBenchmark):
    def __init__(self, data_path: Union[Path, str, None] = None,
                 rng: Union[np.random.RandomState, int, None] = None, **kwargs):

        data_path = self._try_download_api_file(data_path)

        from tabular_benchmarks.nas_cifar10 import NASCifar10A
        benchmark = NASCifar10A(data_dir=str(data_path), multi_fidelity=True)
        super(NASCifar10ABenchmark, self).__init__(benchmark=benchmark, data_path=data_path, rng=rng, **kwargs)

    @staticmethod
    def get_configuration_space(seed: Union[int, None] = None) -> CS.ConfigurationSpace:
        """
        Return the configuration space for the NASCifar10A benchmark.
        Parameters
        ----------
        seed : int, None
            Random seed for the configuration space.

        Returns
        -------
            CS.ConfigurationSpace - Containing the benchmark's hyperparameter
        """

        from tabular_benchmarks.nas_cifar10 import NASCifar10A
        return NASCifar10BBenchmark._get_configuration_space(NASCifar10A, seed)

    def _query_benchmark(self, config: Dict, run_index: int, budget: int = 108) -> Dict:
        """
        Copied from the 'objective_function' from nas_cifar10.py
        We adapted the file in such a way, that the complete result is returned. The original implementation returns
        only the validation error. Now, it can also return the test loss for a given configuration.

        Parameters
        ----------
        config : Dict
        run_index : int
            Specifies the seed to use. Can be one of 0, 1, 2.
        budget : int
            The number of epochs. Must be one of: 4 12 36 108. Otherwise a accuracy of 0 is returned.

        Returns
        -------
        Dict
        """

        failure = {"test_accuracy": 0, "train_accuracy": 0, "validation_accuracy": 0, "training_time": 0,
                   "info": "failure", "trainable_parameters": 0, "module_operations": 0}

        if self.benchmark.multi_fidelity is False:
            assert budget == 108

        matrix = np.zeros([VERTICES, VERTICES], dtype=np.int8)
        idx = np.triu_indices(matrix.shape[0], k=1)
        for i in range(VERTICES * (VERTICES - 1) // 2):
            row = idx[0][i]
            col = idx[1][i]
            matrix[row, col] = config["edge_%d" % i]

        # if not graph_util.is_full_dag(matrix) or graph_util.num_edges(matrix) > MAX_EDGES:
        if graph_util.num_edges(matrix) > MAX_EDGES:
            self.benchmark.record_invalid(config, 1, 1, 0)
            return failure

        labeling = [config["op_node_%d" % i] for i in range(5)]
        labeling = ['input'] + list(labeling) + ['output']
        model_spec = api.ModelSpec(matrix, labeling)

        try:
            data = modified_query(self.benchmark, run_index=run_index, model_spec=model_spec, epochs=budget)
        except api.OutOfDomainError:
            self.benchmark.record_invalid(config, 1, 1, 0)
            return failure

        self.benchmark.record_valid(config, data, model_spec)

        # We dont need this field.
        data.pop('module_adjacency')

        return data

class NASCifar10BBenchmark(NASCifar10BaseBenchmark):
    def __init__(self, data_path: Union[Path, str, None] = None,
                 rng: Union[np.random.RandomState, int, None] = None, **kwargs):

        data_path = self._try_download_api_file(data_path)

        from tabular_benchmarks.nas_cifar10 import NASCifar10B
        benchmark = NASCifar10B(data_dir=str(data_path), multi_fidelity=True)
        super(NASCifar10BBenchmark, self).__init__(benchmark=benchmark, data_path=data_path, rng=rng, **kwargs)

    @staticmethod
    def get_configuration_space(seed: Union[int, None] = None) -> CS.ConfigurationSpace:
        """
        Return the configuration space for the NASCifar10B benchmark.
        Parameters
        ----------
        seed : int, None
            Random seed for the configuration space.

        Returns
        -------
            CS.ConfigurationSpace - Containing the benchmark's hyperparameter
        """

        from tabular_benchmarks.nas_cifar10 import NASCifar10B
        return NASCifar10BBenchmark._get_configuration_space(NASCifar10B, seed)

    def _query_benchmark(self, config: Dict, run_index: int, budget: int = 108) -> Dict:
        """
        Copied from the 'objective_function' from nas_cifar10.py
        We adapted the file in such a way, that the complete result is returned. The original implementation returns
        only the validation error. Now, it can also return the test loss for a given configuration.

        Parameters
        ----------
        config : Dict
        budget : int
            The number of epochs. Must be one of: 4 12 36 108. Otherwise a accuracy of 0 is returned.

        Returns
        -------
        Dict
        """
        failure = {"test_accuracy": 0, "train_accuracy": 0, "validation_accuracy": 0, "training_time": 0,
                   "info": "failure", "trainable_parameters": 0, "module_operations": 0}

        if self.benchmark.multi_fidelity is False:
            assert budget == 108

        bitlist = [0] * (VERTICES * (VERTICES - 1) // 2)
        for i in range(MAX_EDGES):
            bitlist[config["edge_%d" % i]] = 1
        out = 0
        for bit in bitlist:
            out = (out << 1) | bit

        matrix = np.fromfunction(graph_util.gen_is_edge_fn(out),
                                 (VERTICES, VERTICES),
                                 dtype=np.int8)
        # if not graph_util.is_full_dag(matrix) or graph_util.num_edges(matrix) > MAX_EDGES:
        if graph_util.num_edges(matrix) > MAX_EDGES:
            self.benchmark.record_invalid(config, 1, 1, 0)
            return failure

        labeling = [config["op_node_%d" % i] for i in range(5)]
        labeling = ['input'] + list(labeling) + ['output']
        model_spec = api.ModelSpec(matrix, labeling)
        try:
            data = modified_query(self.benchmark, run_index=run_index, model_spec=model_spec, epochs=budget)
        except api.OutOfDomainError:
            self.benchmark.record_invalid(config, 1, 1, 0)
            return failure

        self.benchmark.record_valid(config, data, model_spec)

        # We dont need this field.
        data.pop('module_adjacency')

        return data

class NASCifar10CBenchmark(NASCifar10BaseBenchmark):
    def __init__(self, data_path: Union[Path, str, None] = None,
                 rng: Union[np.random.RandomState, int, None] = None, **kwargs):

        data_path = self._try_download_api_file(data_path)

        from tabular_benchmarks.nas_cifar10 import NASCifar10C
        benchmark = NASCifar10C(data_dir=str(data_path), multi_fidelity=True)
        super(NASCifar10CBenchmark, self).__init__(benchmark=benchmark, data_path=data_path, rng=rng, **kwargs)

    @staticmethod
    def get_configuration_space(seed: Union[int, None] = None) -> CS.ConfigurationSpace:
        """
        Return the configuration space for the NASCifar10C benchmark.
        Parameters
        ----------
        seed : int, None
            Random seed for the configuration space.

        Returns
        -------
            CS.ConfigurationSpace - Containing the benchmark's hyperparameter
        """

        from tabular_benchmarks.nas_cifar10 import NASCifar10C
        return NASCifar10BBenchmark._get_configuration_space(NASCifar10C, seed)

    def _query_benchmark(self, config: Dict, run_index: int, budget: int = 108) -> Dict:
        """
        Copied from the 'objective_function' from nas_cifar10.py
        We adapted the file in such a way, that the complete result is returned. The original implementation returns
        only the validation error. Now, it can also return the test loss for a given configuration.

        Parameters
        ----------
        config : Dict
        budget : int
            The number of epochs. Must be one of: 4 12 36 108. Otherwise a accuracy of 0 is returned.

        Returns
        -------
        Dict
        """
        # Unify the return value to a dictionary.
        failure = {"test_accuracy": 0, "train_accuracy": 0, "validation_accuracy": 0, "training_time": 0,
                   "info": "failure", "trainable_parameters": 0, "module_operations": 0}

        if self.benchmark.multi_fidelity is False:
            assert budget == 108

        edge_prob = []
        for i in range(VERTICES * (VERTICES - 1) // 2):
            edge_prob.append(config["edge_%d" % i])

        idx = np.argsort(edge_prob)[::-1][:config["num_edges"]]
        binay_encoding = np.zeros(len(edge_prob))
        binay_encoding[idx] = 1
        matrix = np.zeros([VERTICES, VERTICES], dtype=np.int8)
        idx = np.triu_indices(matrix.shape[0], k=1)
        for i in range(VERTICES * (VERTICES - 1) // 2):
            row = idx[0][i]
            col = idx[1][i]
            matrix[row, col] = binay_encoding[i]

        if graph_util.num_edges(matrix) > MAX_EDGES:
            self.benchmark.record_invalid(config, 1, 1, 0)
            return failure

        labeling = [config["op_node_%d" % i] for i in range(5)]
        labeling = ['input'] + list(labeling) + ['output']
        model_spec = api.ModelSpec(matrix, labeling)
        try:
            data = modified_query(self.benchmark, run_index=run_index, model_spec=model_spec, epochs=budget)
        except api.OutOfDomainError:
            self.benchmark.record_invalid(config, 1, 1, 0)
            return failure

        self.benchmark.record_valid(config, data, model_spec)

        # We dont need this field.
        data.pop('module_adjacency')

        return data

def modified_query(benchmark, model_spec, run_index: int, epochs=108, stop_halfway=False):
    """
    NOTE:
    Copied from https://github.com/google-research/nasbench/blob/b94247037ee470418a3e56dcb83814e9be83f3a8/nasbench/api.py#L204-L263  # noqa
    We changed the function in such a way that we now can specified the run index (index of the evaluation) which was
    in the original code sampled randomly.

    OLD DOCSTRING:
    Fetch one of the evaluations for this model spec.

    Each call will sample one of the config['num_repeats'] evaluations of the
    model. This means that repeated queries of the same model (or isomorphic
    models) may return identical metrics.

    This function will increment the budget counters for benchmarking purposes.
    See self.training_time_spent, and self.total_epochs_spent.

    This function also allows querying the evaluation metrics at the halfway
    point of training using stop_halfway. Using this option will increment the
    budget counters only up to the halfway point.

    Args:
      model_spec: ModelSpec object.
      epochs: number of epochs trained. Must be one of the evaluated number of
        epochs, [4, 12, 36, 108] for the full dataset.
      stop_halfway: if True, returned dict will only contain the training time
        and accuracies at the halfway point of training (num_epochs/2).
        Otherwise, returns the time and accuracies at the end of training
        (num_epochs).

    Returns:
      dict containing the evaluated data for this object.

    Raises:
      OutOfDomainError: if model_spec or num_epochs is outside the search space.
    """
    if epochs not in benchmark.dataset.valid_epochs:
        raise OutOfDomainError('invalid number of epochs, must be one of %s'
                               % benchmark.dataset.valid_epochs)

    fixed_stat, computed_stat = benchmark.dataset.get_metrics_from_spec(model_spec)

    # MODIFICATION: Use the run index instead of the sampled one.
    # sampled_index = random.randint(0, self.config['num_repeats'] - 1)
    computed_stat = computed_stat[epochs][run_index]

    data = {}
    data['module_adjacency'] = fixed_stat['module_adjacency']
    data['module_operations'] = fixed_stat['module_operations']
    data['trainable_parameters'] = fixed_stat['trainable_parameters']

    if stop_halfway:
        data['training_time'] = computed_stat['halfway_training_time']
        data['train_accuracy'] = computed_stat['halfway_train_accuracy']
        data['validation_accuracy'] = computed_stat['halfway_validation_accuracy']
        data['test_accuracy'] = computed_stat['halfway_test_accuracy']
    else:
        data['training_time'] = computed_stat['final_training_time']
        data['train_accuracy'] = computed_stat['final_train_accuracy']
        data['validation_accuracy'] = computed_stat['final_validation_accuracy']
        data['test_accuracy'] = computed_stat['final_test_accuracy']

    benchmark.dataset.training_time_spent += data['training_time']
    if stop_halfway:
        benchmark.dataset.total_epochs_spent += epochs // 2
    else:
        benchmark.dataset.total_epochs_spent += epochs

    return data


# To add a new cell, type '# %%'
# To add a new markdown cell, type '# %% [markdown]'
# %%
from IPython import get_ipython

# %%
from IPython import get_ipython

# %% [markdown]
#  # COVID-19 Epidemiology Models
# %% [markdown]
#  First, some preliminary imports. You may need to pip install `holoviews` and `GitPython`. For `holoviews`, some extras might be needed (see https://holoviews.org/install.html).

# %%
import pandas as pd
import numpy as np
import scipy.integrate
import matplotlib.pyplot as plt

import bokeh.io
import bokeh.application
import bokeh.application.handlers
import bokeh.models

import holoviews as hv

bokeh.io.output_notebook()
hv.extension('bokeh')

# %% [markdown]
#  Let's load the data from the relevant folder. If this data doesn't exist for you, you'll need to run the `processing/raw_data_processing/daily_refresh.sh` script (which may require `pip install us`).

# %%
import git

repo = git.Repo("./", search_parent_directories=True)
homedir = repo.working_dir
datadir = f"{homedir}/data/us/"

# %% [markdown]
#  Load the US data by state (50 states, but we also make a bunch of other countries pay taxes without getting the right to vote (Puerto Rico, Guam, etc.), so we include their data too).
#  This results in more than 50 "regions".

# %%
df = pd.read_csv(datadir + 'covid/nyt_us_states.csv')

# %% [markdown]
# We want to count days as integers from some starting point.

# %%
df['date_processed'] = pd.to_datetime(df['date'].values)
df['date_processed'] = (df['date_processed'] - df['date_processed'].min()) / np.timedelta64(1, 'D')

# %% [markdown]
# Now we import the population for each state, and add that information correctly to the Pandas "dataframe".

# %%
populations = pd.read_csv(datadir + 'demographics/state_populations_augmented.csv')
def get_population(region):
    return populations[populations['state'] == region]['population'].values[0]

def get_ticker(state):
    state_to_ticker = dict([('CA','California'),('TX','Texas'),('FL','Florida'),('NY','New York'),('PA','Pennsylvania'),('IL','Illinois'),('OH','Ohio'),('GA','Georgia'),('NC','North Carolina'),
                           ('MI','Michigan'),('NJ','New Jersey'),('VA','Virginia'),('WA','Washington'),('AZ','Arizona'),('MA','Massachusetts'),('TN','Tennessee'),('IN','Indiana'),('MO','Missouri'),
                           ('MD','Maryland'),('WI','Wisconsin'),('CO','Colorado'),('MN','Minnesota'),('SC','South Carolina'),('AL','Alabama'),('LA','Louisiana'),('KY','Kentucky'),('OR','Oregon'),
                           ('OK','Oklahoma'),('CT','Connecticut'),('UT','Utah'),('IA','Iowa'),('NV','Nevada'),('AR','Arkansas'),('MS','Mississippi'),('KS','Kansas'),('NM','New Mexico'),('NE','Nebraska'),
                           ('WV','West Virginia'),('ID','Idaho'),('HI','Hawaii'),('NH','New Hampshire'),('ME','Maine'),('MT','Montana'),('RI','Rhode Island'),('DE','Delaware'),('SD','South Dakota'),
                           ('ND','North Dakota'),('AK','Alaska'),('DC','District of Columbia'),('VT','Vermont'),('WY','Wyoming'),('PR','Puerto Rico'),('VI','Virgin Islands'),('GU','Guam'),
                           ('NMI','Northern Mariana Islands')])
    state_to_ticker = {v: k for k, v in state_to_ticker.items()} # I accidentally typed the dictionary backwards (:
    return state_to_ticker[state]

# %% [markdown]
# Let's also include the state ticker (i.e. Minnesota --> MN) as a column in the Pandas dataframe. (Because the US population info is given by state tickers).

# %%
df['state_ticker'] = df.apply(lambda row: get_ticker(row.state),axis = 1)
df['Population'] = df.apply(lambda row: get_population(row.state_ticker), axis=1)

# %% [markdown]
#  Just checking to make sure the data is here...

# %%
df.head()

# %% [markdown]
#  Great! Let's also make a helper function to select data from a state, starting when the pandemic hit to be able to fit models.
#
#  We could change the parameter "min_deaths" here if we want... it seems like an arbitrary choice.

# %%
# return data ever since first min_cases cases
def select_region(df, region, min_deaths=50):
    d = df.loc[df['state'] == region]
    start = np.where(d['deaths'].values > min_deaths)[0][0]
    d = d[start:]
    return d

# %% [markdown]
#  ## MODEL 1 of 3: erf model
# %% [markdown]
#  Let's start with a simple model used in this major paper by IHME: https://www.medrxiv.org/content/10.1101/2020.03.27.20043752v1. Just fit an erf function, nothing else.

# %%
from scipy.special import erf
from scipy.optimize import curve_fit

def erf_curve(t, logp, a, b):
    p = 10**logp
    deaths = p/2*(1+erf(a*(t-b)))
    return deaths

def erf_model(params, data, future=0):
    # initial conditions
    p = 10**params[0]
    a = params[1]
    b = params[2]

    t = data['date_processed'].values
    if future > 0:
        extrapolation = np.arange(future)
        t = np.concatenate((t, extrapolation + t[-1] + 1))
    deaths = p/2*(1+erf(a*(t-b)))

    return t, deaths

def erf_fit(params, data):
    Draw = data['deaths'].values
    Draw[np.where(Draw < 1)] = 1
    Ddata = Draw
    t, Draw = erf_model(params, data)
    Draw[np.where(Draw < 1)] = 1
    D = Draw

    error = mse(D, Ddata)

    return error

# %% [markdown]
#  Now let's try to have something to run the fit and plot the results, including a prediction. To estimate errors, we'll use `curve_fit` and the covariance matrix it returns. Sampling parameters from a Gaussian distribution around the fit, we then bootstrap errors corresponding to 25th and 75th percentile predictions.
#
#  Maybe we can change the way that the quantile predictions are made. The quantile predictions get crazy later (~1e150, which is way to fucking big).

# %%
from bokeh.models import Span

def plot_erf_with_errors_sample(df, region, extrapolate=1, boundary=None):
    '''
    df: main italy data frame
    region: string (region name)
    extrapolate: run model on this * length of data (ex. extrapolate=1 does no prediction, extrapolate=2 predicts a time interval equal in length to the data length)
    boundary: train model on data until the day demarcated by boundary; if boundary is None, train on all available data
    '''
    data = select_region(df, region)
    keys = ['deaths', 'cases']
    all_out = []
    proper = []
    for k in keys:
        erf_params0 = [np.log10(np.max(data[k])), .1, 30]
        y = data[k].values[:boundary]
        popt, pcov = curve_fit(erf_curve, data['date_processed'].values[:boundary], y, p0=erf_params0)
        errors = np.sqrt(np.diag(pcov))
        #======BEGIN=======DEBUGGING================================================================================
        if k == 'deaths':
            print('popt = ', popt)
            print('errors = ', errors)
        #========END=========DEBUGGING================================================================================
        all_s = []
        samples = 100
        for i in range(samples):
            sample = np.random.normal(loc=popt, scale=errors)
            t, y = erf_model(sample, data, future=len(data)*(extrapolate-1))
            all_s.append(y)

        all_s = np.array(all_s)
        all_out.append(all_s)

        t, y = erf_model(popt, data, future=len(data)*(extrapolate-1))
        proper.append(y)

    t = np.arange(0, len(data))
    tp = np.arange(0, len(data)*extrapolate)

    p = bokeh.plotting.figure(plot_width=600,
                              plot_height=400,
                             title = region + ' erf Model',
                             x_axis_label = 't (days)',
                             y_axis_label = '# people')

    p1 = 25
    p2 = 75
    p.varea(x=tp, y1=np.percentile(all_out[1], p1, axis=0), y2=np.percentile(all_out[1], p2, axis=0), color='purple', fill_alpha=0.2)
    p.varea(x=tp, y1=np.percentile(all_out[0], p1, axis=0), y2=np.percentile(all_out[0], p2, axis=0), color='black', fill_alpha=0.2)

#     s1 = np.percentile(all_s, 40, axis=0)
#     s2 = np.percentile(all_s, 60, axis=0)
#     p.varea(x=tp, y1=s1[:, 2], y2=s2[:, 2], color='red', fill_alpha=0.2)
#     p.varea(x=tp, y1=s1[:, 3], y2=s2[:, 3], color='purple', fill_alpha=0.2)
#     p.varea(x=tp, y1=s1[:, 5], y2=s2[:, 5], color='black', fill_alpha=0.2)

    p.line(tp, proper[0], color = 'black', line_width = 1, legend_label = 'Deceased')
    p.line(tp, proper[1] , color = 'purple', line_width = 1, legend_label = 'Symptomatic infected')

    # deaths
    p.circle(t, data[keys[0]], color ='black')

    # cases
    p.circle(t, data[keys[1]], color ='purple')

    if boundary is not None and boundary < len(data):
        vline = Span(location=boundary, dimension='height', line_color='black', line_width=3)
        p.renderers.extend([vline])

    p.legend.location = 'top_left'
    bokeh.io.show(p)

# %% [markdown]
#  And how well does this model do at prediction? We train it on days 0-16 (left of the vertical line), and we try to predict days afterwards.

# %%
plot_erf_with_errors_sample(df, 'Washington', 2, 10)

# %% [markdown]
#  Hm. Let's check how it performs more systematically. Let's try adjusting how long we train and predict for.
#
#  There is obviously some kind of fitting error that's happening here. But I don't care that much because the erf model is not even a great model.

# %%
start = 12
step = 4
ind = 0
results = []
one_more = False
while start + ind*step <= 18:
    boundary = start + ind*step
    plot_erf_with_errors_sample(df, 'Washington', 2, boundary)
    ind += 1

# %% [markdown]
#  The model seems to think that the pandemic will flatten out in the next couple days. Given the earlier plots, I doubt that is right. We can check it on another region and see what happens! What regions can we choose from? We need a region with enough data.

# %%
regions = sorted(np.unique(df['state'].values)) # this is just alphabetical
print(regions)
for r in regions:
    print(r + ' has ',[np.amax(df.loc[df['state'] == r]['cases'].values)], ' cases')

# %% [markdown]
#  The predictions look a little better for Emilia Romagna, but it appears that the error bars for anything earlier than 16-day prediction are enormous. By the time you reach 20 days, however, we only have 22 days of data, so it seems a little trivial to predict.

# %%
# start = 16
# step = 4
# ind = 0
# results = []
# one_more = False
# while start + ind*step <= 24:
#     boundary = start + ind*step
#     plot_erf_with_errors_sample(df, 'Alabama', 2, boundary)
#     ind += 1

# %% [markdown]
#  Time to move on to a more standard epidemiological model.
# %% [markdown]
#  ## MODEL 2 of 3: SEIR-QD model
# %% [markdown]
#  Ok, will this model isn't too standard. It's selected as one of the best-performing models from "Rational evaluation of various epidemic models based on the COVID-19 data of China" (https://www.medrxiv.org/content/10.1101/2020.03.12.20034595v1.full.pdf), where it turned out to be one of the better models in terms of early-stage data collection. It takes into effect quarantine and self-protection, but not from data; it just makes assumptions about how people behave. The differential equations can be found on page 8 of the supplement: https://www.medrxiv.org/content/medrxiv/suppl/2020/03/16/2020.03.12.20034595.DC1/2020.03.12.20034595-1.pdf. Thanks @Liana Merk for doing the initial work on this model!
#
#  Parameters:
#  * $\beta$ = infection rate, from earlier plotting, $10^{(-8)} - 10^{(-6)}$ seem reasonable.
#  * $\delta$ = recovery rate, which we think is on the order of 10-40 days.
#  * $\gamma$ = transition of exposed individuals to infected, which we aren't sure of, especially with the unknown number of asymptomatics.
#  * $\alpha$ = protection rate of susceptible individuals, which we also don't know, and is most likely dynamic over the course of the outbreak.
#  * $\lambda$ = transition rate of infected to quarantined with infection, same as above.
#  * $\kappa$ = death rate, which we think is around 0.01-0.06. We will leave a range between 0.01 and 0.1.

# %%
def seirqd(dat, t, params, N):
    beta = params[0] / N
    delta = params[1]
    gamma = params[2]
    alpha = params[3]
    lambda_ = params[4]
    kappa = params[5]

    s = dat[0]
    e = dat[1]
    i = dat[2]
    q = dat[3]
    r = dat[4]
    d = dat[5]
    sa = dat[6]

    dsdt = - beta * s * i - alpha * s
    dedt = beta * s * i - gamma * e
    didt = gamma * e - lambda_ * i
    dqdt = lambda_ * i - delta * q - kappa * q
    drdt = delta * q
    dddt = kappa * q
    dsadt = alpha * s

    # susceptible, exposed, infected, quarantined, recovered, died, unsusceptible
    return [dsdt, dedt, didt, dqdt, drdt, dddt, dsadt]

# %%
from sklearn.metrics import mean_squared_error

def mse_qd(A, B):
    Ap = np.nan_to_num(A)
    Bp = np.nan_to_num(B)
    Ap[A == -np.inf] = 0
    Bp[B == -np.inf] = 0
    Ap[A == np.inf] = 0
    Bp[B == np.inf] = 0
    return mean_squared_error(Ap, Bp)

def model_qd(params, data, tmax=-1):
    # initial conditions
    N = data['Population'].values[0] # total population
    initial_conditions = N * np.array(params[-5:]) # the parameters are a fraction of the population so multiply by the population

    # initial conditions
    e0 = initial_conditions[0]
    i0 = initial_conditions[1]
    q0 = initial_conditions[2]
    r0 = initial_conditions[3]
    sa0 = initial_conditions[4]

    d0 = data['deaths'].values[0]
    s0 = N - np.sum(initial_conditions) - d0

    yz_0 = np.array([s0, e0, i0, q0, r0, d0, sa0])

    # Package parameters into a tuple
    args = (params, N)

    n = len(data)
    if tmax > 0:
        n = tmax

    # Integrate ODEs
    s = scipy.integrate.odeint(seirqd, yz_0, np.arange(0, n), args=args)

    return s

def fit_leastsq_qd(params, data):
    Ddata = (data['deaths'].values)
    Idata = (data['cases'].values)
    s = model_qd(params, data)

    S = s[:,0]
    E = s[:,1]
    I = s[:,2]
    Q = s[:,3]
    R = s[:,4]
    D = s[:,5]
    SA = s[:,6]

    error = np.concatenate((D-Ddata, I - Idata))
    return error

# %% [markdown]
#  Let's check if this is reasonable with a manual fit to the US state of Washington. Anyway, we need to provide the optimization method an initial guess of the parameters, so this gives us a good chance to make a reasonable guess and speed up the optimization. Some good guesses are provided by the SEIR model here: https://gabgoh.github.io/COVID/index.html.
#
#  NOT SURE WHY THIS FIT DOESN'T WORK WELL

# %%
get_ipython().run_line_magic('matplotlib', 'notebook')
get_ipython().run_line_magic('matplotlib', 'inline')

plt.figure()
d = select_region(df, 'Washington')
# parameters: beta, delta, gamma, alpha, lambda, kappa
params = [2.0, 0.3, 0.2, 0.05, 0.2, 0.03]
# parameters: beta, sigma, ra, shift
# params = [1.1, 0.8, 0.2, 0.3]
# conditions: E, I, Q, R, SA
initial_conditions = [0.5e-3, 0.5e-3, 0.3e-3, 0.1e-4, 0.5]
s = model_qd(params + initial_conditions, d, False)
plt.scatter(d['date_processed'], d['deaths'])
plt.plot(d['date_processed'], s[:, 5])
plt.show()

print('Do the initial conditions sum to one? If not, what do they physically mean?')
print('They sum to ',sum(initial_conditions))

# %% [markdown]
#  Looks good! How about actually fitting it now? We need a plotting function!
#
#  ^ DOES NOT LOOK GOOD! LOOKS BAD! Do we need to change the initial conditions? This model ends up working later, so something must change to make it work for the state of Washington.

# %%
import itertools

def plot_qd(res, p0_params, p0_initial_conditions, df, region, extrapolate=1, boundary=None, plot_infectious=False):
    data = select_region(df, region)

    s = model_qd(res.x, data, len(data)*extrapolate)
    S = s[:,0]
    E = s[:,1]
    I = s[:,2]
    Q = s[:,3]
    R = s[:,4]
    D = s[:,5]
    SA = s[:,6]

    t = np.arange(0, len(data))
    tp = np.arange(0, len(data)*extrapolate)

    p = bokeh.plotting.figure(plot_width=600,
                              plot_height=400,
                             title = region + ' SEIIRD+Q Model',
                             x_axis_label = 't (days)',
                             y_axis_label = '# people')

    if plot_infectious:
        p.line(tp, I, color = 'red', line_width = 1, legend_label = 'All infected')
    p.line(tp, D, color = 'black', line_width = 1, legend_label = 'Deceased')

    # death
    p.circle(t, data['deaths'], color ='black')

    # quarantined
    p.circle(t, data['cases'], color ='purple', legend_label='Tested infected')

    if boundary is not None:
        vline = Span(location=boundary, dimension='height', line_color='black', line_width=3)
        p.renderers.extend([vline])

    p.legend.location = 'top_left'
    bokeh.io.show(p)

# %% [markdown]
#  We'll need to estimate errors from the optimization procedure. I transferred some of the relevant code from `curve_fit` to work for this kind of model:

# %%
from scipy.linalg import svd
# return params, 1 standard deviation errors
def get_errors(res, p0):
    p0 = np.array(p0)
    ysize = len(res.fun)
    cost = 2 * res.cost  # res.cost is half sum of squares!
    popt = res.x
    # Do Moore-Penrose inverse (Pseudoinverse) discarding zero singular values.
    _, s, VT = svd(res.jac, full_matrices=False)
    threshold = np.finfo(float).eps * max(res.jac.shape) * s[0]
    s = s[s > threshold]
    VT = VT[:s.size]
    pcov = np.dot(VT.T / s**2, VT)

    warn_cov = False
    absolute_sigma = False
    if pcov is None:
        # indeterminate covariance
        pcov = zeros((len(popt), len(popt)), dtype=float)
        pcov.fill(inf)
        warn_cov = True
    elif not absolute_sigma:
        if ysize > p0.size:
            s_sq = cost / (ysize - p0.size)
            pcov = pcov * s_sq
        else:
            pcov.fill(inf)
            warn_cov = True

    if warn_cov:
        print('cannot estimate variance')
        return None

    perr = np.sqrt(np.diag(pcov))
    return perr

# %% [markdown]
#  To do the optimization, we should define reasonable ranges for the parameters as well.

# %%
param_ranges = [(0.5, 3.0), (0.0, 0.5), (0.0, 0.5), (0.01, 0.5), (0.0, 0.5), (0.005, 0.1)]
initial_ranges = [(1.0e-7, 0.01), (1.0e-7, 0.01), (1.0e-7, 0.01), (1.0e-7, 0.01), (1.0e-7, 0.9)]
guesses = params + initial_conditions
ranges = param_ranges + initial_ranges

# %%
from scipy.optimize import least_squares
res = least_squares(fit_leastsq_qd, guesses, args=(select_region(df, 'Wisconsin'),), bounds=np.transpose(np.array(ranges)))

# %%
plot_qd(res, params, initial_conditions, df, 'Wisconsin', extrapolate=2, plot_infectious=True)

# %% [markdown]
#  This is quite misleading! We don't actually know how many people are infected, so this model isn't quite what we're looking for. Instead, we only know the people who were tested as infected; they usually have symptoms. I won't bother with the error analysis for this model, but we can see how it holds up to prediction.
#
#  HERE THE MODEL FINALLY WORKS FOR US DATA (Washington). Note that initial conditions are given here. I'm not sure where they got the initial conditions, but in the code, they are defined 5 or 6 cells above.

# %%
start = 10
step = 4
ind = 0
results = []
one_more = False
while start + ind*step <= 18:
    boundary = start + ind*step
    res = least_squares(fit_leastsq_qd, guesses, args=(select_region(df, 'New York')[:boundary],), bounds=np.transpose(np.array(ranges)))
    plot_qd(res, params, initial_conditions, df, 'New York', extrapolate=2, boundary=boundary, plot_infectious=True)
    ind += 1

# %% [markdown]
#  Despite being conceptually wrong, this produces much more satisfactory predictions of fatality rates than the erf model. The infectious predictions should be taken with a (very large) grain of salt anyway, since we know it isn't really predicting the infected population.
# %% [markdown]
#  ## MODEL 3 of 3: $\mathbf{SEI_A I_S R}$ with empirical quarantine
# %% [markdown]
#  The motivation for this model is to add two crucial ingredients: quarantine data and asymptomatic cases. For quarantine analysis, we find an effective population size based on what fraction of the population is moving according to https://citymapper.com/cmi/milan. (Since Milan is the capital of Lombardy, we perform the analysis for that region.) To make the quarantine more realistic, we model a "leaky" quarantine, where the susceptible population is given by the mobility from above plus some offset. To treat asymptomatic cases, we introduce states $I_A$ (asymptomatic) and $I_S$ (symptomatic) according to the following sketch and differential equations:
# %% [markdown]
#  ![SEIIR + quarantine](images/overview.png)
# %% [markdown]
#  Since this is prototyping the model, we manually enter the chart above (raw data is at the link above) and implement it. We also have a testing function $T(t)$ (called `tau(t)` in the code) that allows us to try out different testing strategies for asymptomatic populations. Sorry about the confusing variable names below: parameters are renamed as $\sigma\to$ `alpha`, $s\to$ `sigma`, and $d\to$ `delta`. Not shown in the equations above but included in the diagram is a fixed offset (`offset`) for the leaky quarantine model.
#
#  ***Regarding "TODO fix data imputation", we need to somehow get this data based on US states. Currently, the mobility data (and therefore quarantine data) is hard coded into the vector 'moving'. I'm not sure what shift is. Some arbitrary parameter that probably also needs to be fit...

# %%
# TODO fix data imputation
def q(t, N, shift):
    moving = np.array([57, 54, 52, 51, 49, 47, 46, 45, 44, 43, 39, 37, 34, 23, 19, 13, 10, 7, 6, 5, 5, 7, 6, 5, 4, 4, 3, 3, 4, 4, 3, 3, 3, 3, 2])/100
    q = N*(1-moving) - shift*N
    if np.round(t) >= len(q):
        return q[-1]
    return q[int(np.round(t))]

def tau(t):
    return 0

def seiirq(dat, t, params, N, max_t, offset):
    if t >= max_t:
        return [0]*8
    beta = params[0]
    alpha = params[1] # rate from e to ia
    sigma = params[2] # rate of asymptomatic people becoming symptotic
    ra = params[3] # rate of asymptomatic recovery
    rs = params[4] # rate of symptomatic recovery
    delta = params[5] # death rate
    shift = params[6] # shift quarantine rate vertically from CityMapper data

    s = dat[0]
    e = dat[1]
    i_a = dat[2]
    i_s = dat[3]

    Qind = (q(t + offset, N, shift) - tau(t + offset)*i_a)/(s + e + i_a - tau(t + offset)*i_a)
    Qia = Qind + (1-Qind)*tau(t + offset)

    dsdt = - beta * s * i_a * (1 - Qind) * (1 - Qia) / N
    dedt = beta * s * i_a* (1 - Qind) * (1 - Qia) / N  - alpha * e
    diadt = alpha * e - (sigma + ra) * i_a
    disdt = sigma * i_a - (delta + rs) * i_s
    dddt = delta * i_s
    drdt = ra * i_a + rs * i_s

    # susceptible, exposed, infected, quarantined, recovered, died, unsusceptible
    out = [dsdt, dedt, diadt, disdt, drdt, dddt]
    return out

# %%
from sklearn.metrics import mean_squared_error

def mse(A, B):
    Ap = np.nan_to_num(A)
    Bp = np.nan_to_num(B)
    Ap[A == -np.inf] = 0
    Bp[B == -np.inf] = 0
    Ap[A == np.inf] = 0
    Bp[B == np.inf] = 0
    return mean_squared_error(Ap, Bp)

def model_z(params, data, tmax=-1):
    # initial conditions
    N = data['Population'].values[0] # total population
    initial_conditions = N * np.array(params[-4:]) # the parameters are a fraction of the population so multiply by the population

    e0 = initial_conditions[0]
    ia0 = initial_conditions[1]
    is0 = initial_conditions[2]
    r0 = initial_conditions[3]

    d0 = data['deaths'].values[0]
    s0 = N - np.sum(initial_conditions) - d0

    offset = data['date_processed'].min()
    yz_0 = np.array([s0, e0, ia0, is0, r0, d0])

    n = len(data)
    if tmax > 0:
        n = tmax

    # Package parameters into a tuple
    args = (params, N, n, offset)

    # Integrate ODEs
    try:
        s = scipy.integrate.odeint(seiirq, yz_0, np.arange(0, n), args=args)
    except RuntimeError:
#         print('RuntimeError', params)
        return np.zeros((n, len(yz_0)))

    return s

def fit_leastsq_z(params, data):
    Ddata = (data['deaths'].values)
    Idata = (data['cases'].values)
    s = model_z(params, data)

    S = s[:,0]
    E = s[:,1]
    I_A = s[:,2]
    I_S = s[:,3]
    R = s[:,4]
    D = s[:,5]

    error = np.concatenate(((D-Ddata), I_S - Idata))
    return error

# %% [markdown]
#  Again, we find some good initial parameters.
#
#  *** We need to find some good initial parameters. ):

# %%
get_ipython().run_line_magic('matplotlib', 'notebook')
get_ipython().run_line_magic('matplotlib', 'inline')

plt.figure()
d = select_region(df, 'Washington')
# parameters: beta, alpha, sigma, ra, rs, delta, shift
params = [1.8, 0.35, 0.1, 0.15, 0.34, 0.015, 0.5]
# conditions: E, IA, IS, R
initial_conditions = [4e-6, 0.0009, 0.0005, 0.0002]
s = model_z(params + initial_conditions, d)
plt.scatter(d['date_processed'], d['deaths'])
plt.plot(d['date_processed'], s[:, 5])
plt.show()

# %%
import itertools

def plot_with_errors_sample_z(res, p0_params, p0_initial_conditions, df, region, extrapolate=1, boundary=None, plot_infectious=False):
    data = select_region(df, region)
    errors = get_errors(res, list(p0_params) + list(p0_initial_conditions))
    errors[len(p0_params):] = 0

    all_s = []
    samples = 100
    for i in range(samples):
        sample = np.random.normal(loc=res.x, scale=errors)
        s = model_z(sample, data, len(data)*extrapolate)
        all_s.append(s)

    all_s = np.array(all_s)

    s = model_z(res.x, data, len(data)*extrapolate)
    S = s[:,0]
    E = s[:,1]
    I_A = s[:,2]
    I_S = s[:,3]
    R = s[:,4]
    D = s[:,5]

    t = np.arange(0, len(data))
    tp = np.arange(0, len(data)*extrapolate)

    p = bokeh.plotting.figure(plot_width=600,
                              plot_height=400,
                             title = region + ' SEIIRD+Q Model',
                             x_axis_label = 't (days)',
                             y_axis_label = '# people')

    s1 = np.percentile(all_s, 25, axis=0)
    s2 = np.percentile(all_s, 75, axis=0)
    if plot_infectious:
        p.varea(x=tp, y1=s1[:, 2], y2=s2[:, 2], color='red', fill_alpha=0.2)
    p.varea(x=tp, y1=s1[:, 3], y2=s2[:, 3], color='purple', fill_alpha=0.2)
    p.varea(x=tp, y1=s1[:, 5], y2=s2[:, 5], color='black', fill_alpha=0.2)

    if plot_infectious:
        p.line(tp, I_A, color = 'red', line_width = 1, legend_label = 'Asymptomatic infected')
    p.line(tp, D, color = 'black', line_width = 1, legend_label = 'Deceased')
    p.line(tp, I_S , color = 'purple', line_width = 1, legend_label = 'Symptomatic infected')

    # death
    p.circle(t, data['deaths'], color ='black')

    # quarantined
    p.circle(t, data['cases'], color ='purple')

    if boundary is not None:
        vline = Span(location=boundary, dimension='height', line_color='black', line_width=3)
        p.renderers.extend([vline])

    p.legend.location = 'top_left'
    bokeh.io.show(p)
    return all_s,s1,s2

# %% [markdown]
#  Let's define the initial ranges of the constants for the ODE.

# %%
# beta, alpha, sigma, ra, rs, delta, shift
param_ranges = [(1.0, 2.0), (0.1, 0.5), (0.1, 0.5), (0.05, 0.5), (0.32, 0.36), (0.005, 0.05), (0.1, 0.6)]
initial_ranges = [(1.0e-7, 0.001), (1.0e-7, 0.001), (1.0e-7, 0.001), (1.0e-7, 0.001)]

guesses = params + initial_conditions
ranges = param_ranges + initial_ranges

# %%
# beta, alpha, sigma, ra, rs, delta, shift
param_ranges = [(1.0, 2.0), (0.1, 0.5), (0.1, 0.5), (0.05, 0.5), (0.32, 0.36), (0.005, 0.05), (0.1, 0.6)]
initial_ranges = [(1.0e-7, 0.001), (1.0e-7, 0.001), (1.0e-7, 0.001), (1.0e-7, 0.001)]

guesses = params + initial_conditions
ranges = param_ranges + initial_ranges

boundary = 16
res = least_squares(fit_leastsq_z, guesses, args=(select_region(df, 'Washington')[:boundary],), bounds=np.transpose(np.array(ranges)),jac = '3-point')
all_s,s1,s2 = plot_with_errors_sample_z(res, params, initial_conditions, df, 'Washington', extrapolate=1, boundary=boundary, plot_infectious=True)
perr = get_errors(res,np.zeros((11,1)))
print('parameter standard deviations = ',perr)

# %%
start = 8
step = 4
ind = 0
results = []
one_more = False
while start + ind*step <= 18:
    boundary = start + ind*step
    res = least_squares(fit_leastsq_z, guesses, args=(select_region(df, 'Washington')[:boundary],), bounds=np.transpose(np.array(ranges)),jac = '3-point')
    plot_with_errors_sample_z(res, params, initial_conditions, df, 'Washington', extrapolate=1, boundary=boundary, plot_infectious=True)
    ind += 1

# %% [markdown]
# ***HUUUUGE uncertainty bars. Need to figure out why. Maybe because we don't have good fit parameters, so the covariance is huge.
#
#
#  Just like SEIR-QD, the predictions of cases aren't so great, although the predictions of fatalaties (which matters more and has less data bias) is reasonably accurate. The model also produces prediction that around 2/3 of the cases are asymptomatic (or at least not tested). This corresponds roughly to some recent studies, such as the 50-75% number reported after testing an entire town of 3,300 in Italy (https://www.repubblica.it/salute/medicina-e-ricerca/2020/03/16/news/coronavirus_studio_il_50-75_dei_casi_a_vo_sono_asintomatici_e_molto_contagiosi-251474302/?ref=RHPPTP-BH-I251454518-C12-P3-S2.4-T1) and similar results in Iceland (https://www.government.is/news/article/2020/03/15/Large-scale-testing-of-general-population-in-Iceland-underway/).

# %%


from .insertion_sort import insertion_sort
import numpy as np

def test_sorted_array_returns_same_sorted_array():
    sorted_list = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
    assert insertion_sort(sorted_list) == sorted_list

def test_backward_sorted_array_returns_sorted_array():
    sorted_list = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
    backward_list = [10, 9, 8, 7, 6, 5, 4, 3, 2, 1]
    assert insertion_sort(backward_list) == sorted_list

def test_mixed_array_returns_sorted_array():
    sorted_list = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
    mixed_list = [7, 5, 2, 4, 3, 8, 9, 10, 1, 6]
    assert insertion_sort(mixed_list) == sorted_list

def test_mixed_string_array_returns_sorted_array():
    sorted_list = ['A', 'B', 'C', 'D', 'E', 'F', 'G']
    mixed_list = ['G', 'C', 'A', 'B', 'F', 'E', 'D']
    assert insertion_sort(mixed_list) == sorted_list

def test_modifies_in_place():
    rand = [i for i in np.random.randint(1, 21, 20)]
    assert insertion_sort(rand) == rand

def test_empty_array():
    assert insertion_sort([]) == []

def test_single_element_array():
    assert insertion_sort([1]) == [1]


#!/usr/bin/env python
# coding: utf-8
# # 01-Quantopian Research Basics
# **Please remember that this notebook will only work on Quantopian! Make an account and upload this notebook file.
# These commands and functions won't work except on the Quantopian trading platform!**
# Note a lot of the written markdown text in this notebook comes direclty from the Quantopian docs and tutorials, definitely check those out as well, they're great!
# ## Research
# The notebook format allows us to easily gather information about variuos securities all within the Quantopian platform.
# Keep in mind this is different than the base coding platform of quantopian, which focuses on actually implementing and backtesting trading strategies.
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
# NO NEED TO DO MAGIC INLINE COMMAND ON QUANTOPIAN!
# ## Getting Information
# Let's go over a few key functions:
# * get_pricing()
# * symbols()
# * local_csv()
# * get_backtest()
# * get_fundamentals()
# ## get_pricing()
# The `get_pricing` function provides access to 12 years of US Equity pricing data: the same data used by the Quantopian backtester.
# `get_pricing` returns a <b>pandas object</b>. This could be a panel, dataframe or series depending on the input values.
mcdon = get_pricing('MCD',
                    start_date='2017-01-01',
                    end_date = '2017-02-01',
                    frequency='minute')
mcdon.head()
# In[19]:
mcdon.info()
# Can only go about 12 years back
# which is really all you need for algo trading,
# going back further probably is more noise than signal.
mcdon = get_pricing('MCD',
                    start_date='2005-01-01',
                    end_date = '2017-01-01',
                    frequency='daily')
mcdon['close_price'].plot()
# In[35]:
mcdon['close_price'].pct_change(1).hist_df(bins=100, figsize=(6, 4))
# ## symbols()
# By default `symbols` returns the security object for a ticker symbol. Specify a ticker symbol, or list of symbols, as a string and get a list of security objects back.
# - Use `symbol_reference_date` to identify which date you want the symbol back for a particular ticker symbol.
# - Specify how you would like missing results to be handled with `handle_missing`
mcdon_eq_info = symbols('MCD')
type(mcdon_eq_info)
# In[48]:
for key in mcdon_eq_info.to_dict():
    print(key)
    print(mcdon_eq_info.to_dict()[key])
    print('\n')
# ## get_fundamentals()
# The `get_fundamentals` function provides programmatic access to the Quantopian fundamental database.
# Based on data provided by Morningstar, `get_fundamentals` provides over 600 corporate metrics dating back to 2002 (to match Quantopian's pricing data).
# The data used by this research function is the same data used by the `get_fundamentals` function used in the Quantopian IDE.
# The fields are described in the Quantopian help documents: http://www.quantopian.com/help/fundamentals.
# Have to do this first in the notebook:
fundamentals = init_fundamentals()
# The get_fundamentals() function takes in a SQLAlchemy query which can be quite complicated and strange looking at first.
# Basically it allows you to filter by a variety of fundamentals (things like Market Cap, P/E Ratio, or even city of HQ). Check out the link above for all the things you can filter by!
# Let's walk through a few query examples.
# First call fundamentals and use tab to check out the various options:
fundamentals. # call tab here as in the video!
# Market Cap
my_query = query(fundamentals.valuation.market_cap)
my_funds = get_fundamentals(my_query,'2017-01-01')
my_funds.info()
# Basically just returns the market cap of everything
# for 2017-01-01
my_funds.head()
# What you usualy do is filter by other qualities after the query!
# Only get companies worth 500 billion or more (that's a lot of dough!)
big_companies = (query(fundamentals.valuation.market_cap).
                 filter(fundamentals.valuation.market_cap > 500000000000) )
my_big_funds = get_fundamentals(big_companies,'2017-07-19')
# On
my_big_funds
# In[70]:
7.82 * 10**11
# In[50]:
get_fundamentals()


import numpy as np

# Loading and existing file
arr = np.loadtxt('somefile.txt')

# Saving a new file
np.savetxt('somenewfile.txt', arr)

# Opening an existing file with the append option
f = open('existingfile.txt', 'a')

# Creating some random data to append to the existing file
data2append = np.random.rand(100)

# With np.savetxt we replace the file name with the file handle.
np.savetxt(f, data2append)

f.close()


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import math
import os
import random
import zipfile

import numpy as np
from six.moves import urllib
from six.moves import xrange
import tensorflow as tf

# 第一步: 在下面这个地址下载语料库
url = 'http://mattmahoney.net/dc/'

def data_path(file_name):
    current_dir = os.path.dirname(__file__)
    parent_dir = os.path.abspath(os.path.join(current_dir, os.pardir, os.pardir))
    data_dir = os.path.join(parent_dir, "data_source", file_name)
    return data_dir

def maybe_download(filename, expected_bytes):
    """
    这个函数的功能是：
    如果filename不存在，就在上面的地址下载它。
    如果filename存在，就跳过下载。
    最终会检查文字的字节数是否和expected_bytes相同。
    """
    # if not os.path.exists(filename):
    if not os.path.exists(data_path(filename)):
        print('start downloading...')
        filename, _ = urllib.request.urlretrieve(url + filename, filename)
    statinfo = os.stat(filename)
    if statinfo.st_size == expected_bytes:
        print('Found and verified', filename)
    else:
        print(statinfo.st_size)
        raise Exception(
            'Failed to verify ' + filename + '. Can you get to it with a browser?')
    return filename

# 下载语料库text8.zip并验证下载
# filename = maybe_download('text8.zip', 31344016)
# print(filename)
# print("%%%%%%%%%%%%%%%%%")

filename = data_path("text8.zip")
print(filename)

# 将语料库解压，并转换成一个word的list
def read_data(filename):
    """
    这个函数的功能是：
    将下载好的zip文件解压并读取为word的list
    """
    with zipfile.ZipFile(filename) as f:
        data = tf.compat.as_str(f.read(f.namelist()[0])).split()
    return data

vocabulary = read_data(filename)
print('Data size', len(vocabulary))  # 总长度为1700万左右
# 输出前100个词。
print(vocabulary[0:100])

# 第二步: 制作一个词表，将不常见的词变成一个UNK标识符
# 词表的大小为5万（即我们只考虑最常出现的5万个词）
vocabulary_size = 50000

def build_dataset(words, n_words):
    """
    函数功能：将原始的单词表示变成index
    """
    count = [['UNK', -1]]
    count.extend(collections.Counter(words).most_common(n_words - 1))
    # count = [['UNK', -1], ('the', 1061396), ('of', 593677), ('and', 416629) ...]
    dictionary = dict()
    for word, _ in count:
        dictionary[word] = len(dictionary)
    data = list()
    unk_count = 0
    for word in words:
        if word in dictionary:
            index = dictionary[word]
        else:
            index = 0
            unk_count += 1
        data.append(index)
    count[0][1] = unk_count
    reversed_dictionary = dict(zip(dictionary.values(), dictionary.keys()))
    return data, count, dictionary, reversed_dictionary

data, count, dictionary, reverse_dictionary = build_dataset(vocabulary, vocabulary_size)
del vocabulary  # 删除已节省内存
# 输出最常出现的5个单词
print('Most common words (+UNK)', count[:5])
# 输出转换后的数据库data，和原来的单词（前10个）
print('Sample data', data[:10], [reverse_dictionary[i] for i in data[:10]])
# 我们下面就使用data来制作训练集
print("+++++++++++++++++")
data_index = 0

# 第三步：定义一个函数，用于生成cbow模型用的batch
def generate_batch(batch_size, cbow_window):
    global data_index
    assert cbow_window % 2 == 1
    span = 2 * cbow_window + 1
    # 去除中心word: span - 1
    batch = np.ndarray(shape=(batch_size, span - 1), dtype=np.int32)
    labels = np.ndarray(shape=(batch_size, 1), dtype=np.int32)

    buffer = collections.deque(maxlen=span)
    for _ in range(span):
        buffer.append(data[data_index])
        # 循环选取 data中数据，到尾部则从头开始
        data_index = (data_index + 1) % len(data)

    for i in range(batch_size):
        # target at the center of span
        target = cbow_window
        # 仅仅需要知道context(word)而不需要word
        target_to_avoid = [cbow_window]

        col_idx = 0
        for j in range(span):
            # 略过中心元素 word
            if j == span // 2:
                continue
            batch[i, col_idx] = buffer[j]
            col_idx += 1
        labels[i, 0] = buffer[target]
        # 更新 buffer
        buffer.append(data[data_index])
        data_index = (data_index + 1) % len(data)

    return batch, labels

num_steps = 100001

if __name__ == '__main__':
    batch_size = 128
    embedding_size = 128  # Dimension of the embedding vector.
    cbow_window = 1  # How many words to consider left and right.
    num_skips = 2  # How many times to reuse an input to generate a label.
    # We pick a random validation set to sample nearest neighbors. here we limit the
    # validation samples to the words that have a low numeric ID, which by
    # construction are also the most frequent.
    valid_size = 16  # Random set of words to evaluate similarity on.
    valid_window = 100  # Only pick dev samples in the head of the distribution.
    # pick 16 samples from 100
    valid_examples = np.array(random.sample(range(valid_window), valid_size // 2))
    valid_examples = np.append(valid_examples, random.sample(range(1000, 1000 + valid_window), valid_size // 2))
    num_sampled = 64  # Number of negative examples to sample.

    graph = tf.Graph()

    with graph.as_default(), tf.device('/cpu:0'):

        # Input data.
        train_dataset = tf.placeholder(tf.int32, shape=[batch_size, 2 * cbow_window])
        train_labels = tf.placeholder(tf.int32, shape=[batch_size, 1])
        valid_dataset = tf.constant(valid_examples, dtype=tf.int32)

        # Variables.
        # embedding, vector for each word in the vocabulary
        embeddings = tf.Variable(tf.random_uniform([vocabulary_size, embedding_size], -1.0, 1.0))
        nce_weights = tf.Variable(tf.truncated_normal([vocabulary_size, embedding_size],
                                                      stddev=1.0 / math.sqrt(embedding_size)))
        nce_biases = tf.Variable(tf.zeros([vocabulary_size]))

        # Model.
        # Look up embeddings for inputs.
        # this might efficiently find the embeddings for given ids (traind dataset)
        # manually doing this might not be efficient given there are 50000 entries in embeddings
        embeds = None
        for i in range(2 * cbow_window):
            embedding_i = tf.nn.embedding_lookup(embeddings, train_dataset[:, i])
            print('embedding %d shape: %s' % (i, embedding_i.get_shape().as_list()))
            emb_x, emb_y = embedding_i.get_shape().as_list()
            if embeds is None:
                embeds = tf.reshape(embedding_i, [emb_x, emb_y, 1])
            else:
                embeds = tf.concat([embeds, tf.reshape(embedding_i, [emb_x, emb_y, 1])], 2)

        assert embeds.get_shape().as_list()[2] == 2 * cbow_window
        print("Concat embedding size: %s" % embeds.get_shape().as_list())
        avg_embed = tf.reduce_mean(embeds, 2, keep_dims=False)
        print("Avg embedding size: %s" % avg_embed.get_shape().as_list())

        loss = tf.reduce_mean(tf.nn.nce_loss(nce_weights, nce_biases,
                                             labels=train_labels,
                                             inputs=avg_embed,
                                             num_sampled=num_sampled,
                                             num_classes=vocabulary_size))

        # Optimizer.
        # Note: The optimizer will optimize the softmax_weights AND the embeddings.
        # This is because the embeddings are defined as a variable quantity and the
        # optimizer's `minimize` method will by default modify all variable quantities
        # that contribute to the tensor it is passed.
        # See docs on `tf.train.Optimizer.minimize()` for more details.
        # Adagrad is required because there are too many things to optimize
        optimizer = tf.train.AdagradOptimizer(1.0).minimize(loss)

        # Compute the similarity between minibatch examples and all embeddings.
        # We use the cosine distance:
        norm = tf.sqrt(tf.reduce_sum(tf.square(embeddings), 1, keep_dims=True))
        normalized_embeddings = embeddings / norm
        valid_embeddings = tf.nn.embedding_lookup(normalized_embeddings, valid_dataset)
        similarity = tf.matmul(valid_embeddings, tf.transpose(normalized_embeddings))

    with tf.Session(graph=graph) as session:
        tf.global_variables_initializer().run()
        print('Initialized')
        average_loss = 0
        for step in range(num_steps):
            batch_data, batch_labels = generate_batch(batch_size, cbow_window)
            feed_dict = {train_dataset: batch_data, train_labels: batch_labels}
            _, l = session.run([optimizer, loss], feed_dict=feed_dict)
            average_loss += l
            if step % 2000 == 0:
                if step > 0:
                    average_loss = average_loss / 2000
                    # The average loss is an estimate of the loss over the last 2000 batches.
                print('Average loss at step %d: %f' % (step, average_loss))
                average_loss = 0
            # note that this is expensive (~20% slowdown if computed every 500 steps)
            if step % 10000 == 0:
                sim = similarity.eval()
                for i in range(valid_size):
                    valid_word = reverse_dictionary[valid_examples[i]]
                    top_k = 8  # number of nearest neighbors
                    nearest = (-sim[i, :]).argsort()[1:top_k + 1]
                    log = 'Nearest to %s:' % valid_word
                    for k in range(top_k):
                        close_word = reverse_dictionary[nearest[k]]
                        log = '%s %s,' % (log, close_word)
                    print(log)
        final_embeddings = normalized_embeddings.eval()

# Step 6: 可视化
# 可视化的图片会保存为“tsne1.png”

def plot_with_labels(low_dim_embs, labels, filename='tsne1.png'):
    assert low_dim_embs.shape[0] >= len(labels), 'More labels than embeddings'
    plt.figure(figsize=(18, 18))  # in inches
    for i, label in enumerate(labels):
        x, y = low_dim_embs[i, :]
        plt.scatter(x, y)
        plt.annotate(label,
                     xy=(x, y),
                     xytext=(5, 2),
                     textcoords='offset points',
                     ha='right',
                     va='bottom')

    plt.savefig(filename)

try:
    # pylint: disable=g-import-not-at-top
    from sklearn.manifold import TSNE
    import matplotlib

    matplotlib.use('agg')
    import matplotlib.pyplot as plt

    # 因为我们的embedding的大小为128维，没有办法直接可视化
    # 所以我们用t-SNE方法进行降维
    tsne = TSNE(perplexity=30, n_components=2, init='pca', n_iter=5000)
    # 只画出500个词的位置
    plot_only = 500
    low_dim_embs = tsne.fit_transform(final_embeddings[:plot_only, :])
    labels = [reverse_dictionary[i] for i in xrange(plot_only)]
    plot_with_labels(low_dim_embs, labels)

except ImportError:
    print('Please install sklearn, matplotlib, and scipy to show embeddings.')


# -*- coding: utf-8 -*-
# author: ysoftman
# python version : 3.x
# desc : pandas test
import numpy as np
import pandas as pd

fruits = ["lemon", "orange", "apple"]
print("fruits = ", fruits)
# series(시리즈)는 키를 가지는 리스트의 1차원자료 구조다.
# 시리즈의를 묶어 2차원의 dataframe 구조를 만들 수 있다.
print("pd.Series(fruits) =\n", pd.Series(fruits), sep="")
print()

# int 리스트를 가지고 있는 시리즈
numbers = [1, 2, 3]
print("numbers = ", numbers)
print("pd.Series(numbers) =\n", pd.Series(numbers), sep="")
print()

# None 은 NaN(Not A Number:표현 불가능한 수치)로 표현된다.
# NaN 으로 float 값으로 처리된다.
floats = [1, 2, None]
print("floats = ", floats)
print("pd.Series(floats) =\n", pd.Series(floats), sep="")
print("pd.Series(floats).name =\n", pd.Series(floats).name, sep="")
print()

# 시리즈에 추가하기
fruits1 = pd.Series(fruits, name="new1")
fruits2 = pd.Series(["strawberry", "watermelon", "mango"], name="new2")
fruits3 = pd.Series(["aa", "cc"], name="new3")
fruits1 = fruits1.append(fruits2)
fruits1 = fruits1.append(fruits3)
print("fruits1 =\n", fruits1, sep="")
print("fruits1[1] =\n", fruits1[1], sep="")
print()

# dictionary -> series
fruitsDict = {"lemon": 9000, "apple": "7000", "orange": "5000"}
print("fruitsDict = ", fruitsDict)
print("pd.Series(fruitsDict) =\n", pd.Series(fruitsDict), sep="")
print()

fs = pd.Series(fruitsDict)
# 현재 index
print("fs.index = ", fs.index)
# index 새로 지정할 수 있고
# index 를 지정하지 않는 값은 사라지고(무효화)  => orange
# index 에 해당하는 값이 없으면 NaN 또는 None 이된다. => banana
fs2 = pd.Series(fs, index=["lemon", "apple", "banana"])
print("fs =\n", fs, sep="")
print("fs2 = \n", fs2, sep="")
print()

# label 로 시리즈 인덱싱, 슬라이싱
print('fs.loc["apple"]', fs.loc["apple"])
print('fs.loc["lemon":"orange"]', fs.loc["lemon":"orange"])
# 위치 로 시리즈 인덱싱, 슬라이싱
print('fs.iloc[1]', fs.iloc[1])
print('fs.iloc[0:]', fs.iloc[0:])
# 일반 인덱싱으로도 되지만 2차원이상의 데이터에서는 키 에러가 발생한다.
# print('fs["apple"]', fs["apple"])
# print('fs[1]', fs[1])
# 값 추가
# fs['id'] = 'ysfotman'
fs.loc['id'] = 'ysfotman'
print("fs =\n", fs, sep="")
print()

# 10000 개의 데이터를 담은 시리즈 생성
rs = pd.Series(np.random.randint(1, 100, 10000))
print("len(rs) = ", len(rs))
# 앞부분과 뒤부분 데이터 보기
print("rs.head() =\n", rs.head(n=10), sep="")
print("rs.tail() =\n", rs.tail(), sep="")
print()

# rs 데이터에 일괄 연산 적용
# 다음 처럼 루프를 돌리는것 보다
for k, v in rs.iteritems():
    rs.loc[k] = v + 2
print("rs.tail() =\n", rs.tail(), sep="")
# 벡터화 연산(vectorization operation:루프없이 연산 한번에 끝낸다.)을 지원해 더 빠르다.
rs += 10
print("rs.tail() =\n", rs.tail(), sep="")


import numpy as np
import pandas as pd

column_names = ['color', 'director_name', 'num_critic_for_reviews', 'duration',
                'gross', 'movie_title', 'num_user_for_reviews',	'country',
                'cotent_rating', 'budget', 'title_year', 'imdb_score', 'genre']


# -*- coding: utf-8 -*-
"""03.07-Merge-and-Join.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/github/jakevdp/PythonDataScienceHandbook/blob/master/notebooks/03.07-Merge-and-Join.ipynb

*This notebook contains an excerpt from the [Python Data Science Handbook](http://shop.oreilly.com/product/0636920034919.do) by Jake VanderPlas; the content is available [on GitHub](https://github.com/jakevdp/PythonDataScienceHandbook).*

*The text is released under the [CC-BY-NC-ND license](https://creativecommons.org/licenses/by-nc-nd/3.0/us/legalcode), and code is released under the [MIT license](https://opensource.org/licenses/MIT). If you find this content useful, please consider supporting the work by [buying the book](http://shop.oreilly.com/product/0636920034919.do)!*

< [Combining Datasets: Concat and Append](03.06-Concat-And-Append.ipynb) | [Contents](Index.ipynb) | [Aggregation and Grouping](03.08-Aggregation-and-Grouping.ipynb) >

# Combining Datasets: Merge and Join

One essential feature offered by Pandas is its high-performance, in-memory join and merge operations.
If you have ever worked with databases, you should be familiar with this type of data interaction.
The main interface for this is the ``pd.merge`` function, and we'll see few examples of how this can work in practice.

For convenience, we will start by redefining the ``display()`` functionality from the previous section:
"""

import pandas as pd
import numpy as np

class display(object):
    """Display HTML representation of multiple objects"""
    template = """

{0}
{1}

"""
    def __init__(self, *args):
        self.args = args

    def _repr_html_(self):
        return '\n'.join(self.template.format(a, eval(a)._repr_html_())
                         for a in self.args)

    def __repr__(self):
        return '\n\n'.join(a + '\n' + repr(eval(a))
                           for a in self.args)

"""## Relational Algebra

The behavior implemented in ``pd.merge()`` is a subset of what is known as *relational algebra*, which is a formal set of rules for manipulating relational data, and forms the conceptual foundation of operations available in most databases.
The strength of the relational algebra approach is that it proposes several primitive operations, which become the building blocks of more complicated operations on any dataset.
With this lexicon of fundamental operations implemented efficiently in a database or other program, a wide range of fairly complicated composite operations can be performed.

Pandas implements several of these fundamental building-blocks in the ``pd.merge()`` function and the related ``join()`` method of ``Series`` and ``Dataframe``s.
As we will see, these let you efficiently link data from different sources.

## Categories of Joins

The ``pd.merge()`` function implements a number of types of joins: the *one-to-one*, *many-to-one*, and *many-to-many* joins.
All three types of joins are accessed via an identical call to the ``pd.merge()`` interface; the type of join performed depends on the form of the input data.
Here we will show simple examples of the three types of merges, and discuss detailed options further below.

### One-to-one joins

Perhaps the simplest type of merge expresion is the one-to-one join, which is in many ways very similar to the column-wise concatenation seen in [Combining Datasets: Concat & Append](03.06-Concat-And-Append.ipynb).
As a concrete example, consider the following two ``DataFrames`` which contain information on several employees in a company:
"""

df1 = pd.DataFrame({'employee': ['Bob', 'Jake', 'Lisa', 'Sue'],
                    'group': ['Accounting', 'Engineering', 'Engineering', 'HR']})
df2 = pd.DataFrame({'employee': ['Lisa', 'Bob', 'Jake', 'Sue'],
                    'hire_date': [2004, 2008, 2012, 2014]})
display('df1', 'df2')

"""To combine this information into a single ``DataFrame``, we can use the ``pd.merge()`` function:"""

df3 = pd.merge(df1, df2)
df3

"""The ``pd.merge()`` function recognizes that each ``DataFrame`` has an "employee" column, and automatically joins using this column as a key.
The result of the merge is a new ``DataFrame`` that combines the information from the two inputs.
Notice that the order of entries in each column is not necessarily maintained: in this case, the order of the "employee" column differs between ``df1`` and ``df2``, and the ``pd.merge()`` function correctly accounts for this.
Additionally, keep in mind that the merge in general discards the index, except in the special case of merges by index (see the ``left_index`` and ``right_index`` keywords, discussed momentarily).

### Many-to-one joins

Many-to-one joins are joins in which one of the two key columns contains duplicate entries.
For the many-to-one case, the resulting ``DataFrame`` will preserve those duplicate entries as appropriate.
Consider the following example of a many-to-one join:
"""

df4 = pd.DataFrame({'group': ['Accounting', 'Engineering', 'HR'],
                    'supervisor': ['Carly', 'Guido', 'Steve']})
display('df3', 'df4', 'pd.merge(df3, df4)')

"""The resulting ``DataFrame`` has an aditional column with the "supervisor" information, where the information is repeated in one or more locations as required by the inputs.

### Many-to-many joins

Many-to-many joins are a bit confusing conceptually, but are nevertheless well defined.
If the key column in both the left and right array contains duplicates, then the result is a many-to-many merge.
This will be perhaps most clear with a concrete example.
Consider the following, where we have a ``DataFrame`` showing one or more skills associated with a particular group.
By performing a many-to-many join, we can recover the skills associated with any individual person:
"""

df5 = pd.DataFrame({'group': ['Accounting', 'Accounting',
                              'Engineering', 'Engineering', 'HR', 'HR'],
                    'skills': ['math', 'spreadsheets', 'coding', 'linux',
                               'spreadsheets', 'organization']})
display('df1', 'df5', "pd.merge(df1, df5)")

"""These three types of joins can be used with other Pandas tools to implement a wide array of functionality.
But in practice, datasets are rarely as clean as the one we're working with here.
In the following section we'll consider some of the options provided by ``pd.merge()`` that enable you to tune how the join operations work.

## Specification of the Merge Key

We've already seen the default behavior of ``pd.merge()``: it looks for one or more matching column names between the two inputs, and uses this as the key.
However, often the column names will not match so nicely, and ``pd.merge()`` provides a variety of options for handling this.

### The ``on`` keyword

Most simply, you can explicitly specify the name of the key column using the ``on`` keyword, which takes a column name or a list of column names:
"""

display('df1', 'df2', "pd.merge(df1, df2, on='employee')")

"""This option works only if both the left and right ``DataFrame``s have the specified column name.

### The ``left_on`` and ``right_on`` keywords

At times you may wish to merge two datasets with different column names; for example, we may have a dataset in which the employee name is labeled as "name" rather than "employee".
In this case, we can use the ``left_on`` and ``right_on`` keywords to specify the two column names:
"""

df3 = pd.DataFrame({'name': ['Bob', 'Jake', 'Lisa', 'Sue'],
                    'salary': [70000, 80000, 120000, 90000]})
display('df1', 'df3', 'pd.merge(df1, df3, left_on="employee", right_on="name")')

"""The result has a redundant column that we can drop if desired–for example, by using the ``drop()`` method of ``DataFrame``s:"""

pd.merge(df1, df3, left_on="employee", right_on="name").drop('name', axis=1)

"""### The ``left_index`` and ``right_index`` keywords

Sometimes, rather than merging on a column, you would instead like to merge on an index.
For example, your data might look like this:
"""

df1a = df1.set_index('employee')
df2a = df2.set_index('employee')
display('df1a', 'df2a')

"""You can use the index as the key for merging by specifying the ``left_index`` and/or ``right_index`` flags in ``pd.merge()``:"""

display('df1a', 'df2a',
        "pd.merge(df1a, df2a, left_index=True, right_index=True)")

"""For convenience, ``DataFrame``s implement the ``join()`` method, which performs a merge that defaults to joining on indices:"""

display('df1a', 'df2a', 'df1a.join(df2a)')

"""If you'd like to mix indices and columns, you can combine ``left_index`` with ``right_on`` or ``left_on`` with ``right_index`` to get the desired behavior:"""

display('df1a', 'df3', "pd.merge(df1a, df3, left_index=True, right_on='name')")

"""All of these options also work with multiple indices and/or multiple columns; the interface for this behavior is very intuitive.
For more information on this, see the ["Merge, Join, and Concatenate" section](http://pandas.pydata.org/pandas-docs/stable/merging.html) of the Pandas documentation.

## Specifying Set Arithmetic for Joins

In all the preceding examples we have glossed over one important consideration in performing a join: the type of set arithmetic used in the join.
This comes up when a value appears in one key column but not the other. Consider this example:
"""

df6 = pd.DataFrame({'name': ['Peter', 'Paul', 'Mary'],
                    'food': ['fish', 'beans', 'bread']},
                   columns=['name', 'food'])
df7 = pd.DataFrame({'name': ['Mary', 'Joseph'],
                    'drink': ['wine', 'beer']},
                   columns=['name', 'drink'])
display('df6', 'df7', 'pd.merge(df6, df7)')

"""Here we have merged two datasets that have only a single "name" entry in common: Mary.
By default, the result contains the *intersection* of the two sets of inputs; this is what is known as an *inner join*.
We can specify this explicitly using the ``how`` keyword, which defaults to ``"inner"``:
"""

pd.merge(df6, df7, how='inner')

"""Other options for the ``how`` keyword are ``'outer'``, ``'left'``, and ``'right'``.
An *outer join* returns a join over the union of the input columns, and fills in all missing values with NAs:
"""

display('df6', 'df7', "pd.merge(df6, df7, how='outer')")

"""The *left join* and *right join* return joins over the left entries and right entries, respectively.
For example:
"""

display('df6', 'df7', "pd.merge(df6, df7, how='left')")

"""The output rows now correspond to the entries in the left input. Using
``how='right'`` works in a similar manner.

All of these options can be applied straightforwardly to any of the preceding join types.

## Overlapping Column Names: The ``suffixes`` Keyword

Finally, you may end up in a case where your two input ``DataFrame``s have conflicting column names.
Consider this example:
"""

df8 = pd.DataFrame({'name': ['Bob', 'Jake', 'Lisa', 'Sue'],
                    'rank': [1, 2, 3, 4]})
df9 = pd.DataFrame({'name': ['Bob', 'Jake', 'Lisa', 'Sue'],
                    'rank': [3, 1, 4, 2]})
display('df8', 'df9', 'pd.merge(df8, df9, on="name")')

"""Because the output would have two conflicting column names, the merge function automatically appends a suffix ``_x`` or ``_y`` to make the output columns unique.
If these defaults are inappropriate, it is possible to specify a custom suffix using the ``suffixes`` keyword:
"""

display('df8', 'df9', 'pd.merge(df8, df9, on="name", suffixes=["_L", "_R"])')

"""These suffixes work in any of the possible join patterns, and work also if there are multiple overlapping columns.

For more information on these patterns, see [Aggregation and Grouping](03.08-Aggregation-and-Grouping.ipynb) where we dive a bit deeper into relational algebra.
Also see the [Pandas "Merge, Join and Concatenate" documentation](http://pandas.pydata.org/pandas-docs/stable/merging.html) for further discussion of these topics.

## Example: US States Data

Merge and join operations come up most often when combining data from different sources.
Here we will consider an example of some data about US states and their populations.
The data files can be found at http://github.com/jakevdp/data-USstates/:
"""

# Following are shell commands to download the data
# !curl -O https://raw.githubusercontent.com/jakevdp/data-USstates/master/state-population.csv
# !curl -O https://raw.githubusercontent.com/jakevdp/data-USstates/master/state-areas.csv
# !curl -O https://raw.githubusercontent.com/jakevdp/data-USstates/master/state-abbrevs.csv

"""Let's take a look at the three datasets, using the Pandas ``read_csv()`` function:"""

pop = pd.read_csv('data/state-population.csv')
areas = pd.read_csv('data/state-areas.csv')
abbrevs = pd.read_csv('data/state-abbrevs.csv')

display('pop.head()', 'areas.head()', 'abbrevs.head()')

"""Given this information, say we want to compute a relatively straightforward result: rank US states and territories by their 2010 population density.
We clearly have the data here to find this result, but we'll have to combine the datasets to find the result.

We'll start with a many-to-one merge that will give us the full state name within the population ``DataFrame``.
We want to merge based on the ``state/region``  column of ``pop``, and the ``abbreviation`` column of ``abbrevs``.
We'll use ``how='outer'`` to make sure no data is thrown away due to mismatched labels.
"""

merged = pd.merge(pop, abbrevs, how='outer',
                  left_on='state/region', right_on='abbreviation')
merged = merged.drop('abbreviation', 1) # drop duplicate info
merged.head()

"""Let's double-check whether there were any mismatches here, which we can do by looking for rows with nulls:"""

merged.isnull().any()

"""Some of the ``population`` info is null; let's figure out which these are!"""

merged[merged['population'].isnull()].head()

"""It appears that all the null population values are from Puerto Rico prior to the year 2000; this is likely due to this data not being available from the original source.

More importantly, we see also that some of the new ``state`` entries are also null, which means that there was no corresponding entry in the ``abbrevs`` key!
Let's figure out which regions lack this match:
"""

merged.loc[merged['state'].isnull(), 'state/region'].unique()

"""We can quickly infer the issue: our population data includes entries for Puerto Rico (PR) and the United States as a whole (USA), while these entries do not appear in the state abbreviation key.
We can fix these quickly by filling in appropriate entries:
"""

merged.loc[merged['state/region'] == 'PR', 'state'] = 'Puerto Rico'
merged.loc[merged['state/region'] == 'USA', 'state'] = 'United States'
merged.isnull().any()

"""No more nulls in the ``state`` column: we're all set!

Now we can merge the result with the area data using a similar procedure.
Examining our results, we will want to join on the ``state`` column in both:
"""

final = pd.merge(merged, areas, on='state', how='left')
final.head()

"""Again, let's check for nulls to see if there were any mismatches:"""

final.isnull().any()

"""There are nulls in the ``area`` column; we can take a look to see which regions were ignored here:"""

final['state'][final['area (sq. mi)'].isnull()].unique()

"""We see that our ``areas`` ``DataFrame`` does not contain the area of the United States as a whole.
We could insert the appropriate value (using the sum of all state areas, for instance), but in this case we'll just drop the null values because the population density of the entire United States is not relevant to our current discussion:
"""

final.dropna(inplace=True)
final.head()

"""Now we have all the data we need. To answer the question of interest, let's first select the portion of the data corresponding with the year 2000, and the total population.
We'll use the ``query()`` function to do this quickly (this requires the ``numexpr`` package to be installed; see [High-Performance Pandas: ``eval()`` and ``query()``](03.12-Performance-Eval-and-Query.ipynb)):
"""

data2010 = final.query("year == 2010 & ages == 'total'")
data2010.head()

"""Now let's compute the population density and display it in order.
We'll start by re-indexing our data on the state, and then compute the result:
"""

data2010.set_index('state', inplace=True)
density = data2010['population'] / data2010['area (sq. mi)']

density.sort_values(ascending=False, inplace=True)
density.head()

"""The result is a ranking of US states plus Washington, DC, and Puerto Rico in order of their 2010 population density, in residents per square mile.
We can see that by far the densest region in this dataset is Washington, DC (i.e., the District of Columbia); among states, the densest is New Jersey.

We can also check the end of the list:
"""

density.tail()

"""We see that the least dense state, by far, is Alaska, averaging slightly over one resident per square mile.

This type of messy data merging is a common task when trying to answer questions using real-world data sources.
I hope that this example has given you an idea of the ways you can combine tools we've covered in order to gain insight from your data!

< [Combining Datasets: Concat and Append](03.06-Concat-And-Append.ipynb) | [Contents](Index.ipynb) | [Aggregation and Grouping](03.08-Aggregation-and-Grouping.ipynb) >

"""


import pytest, torch
import numpy as np
from fastai import core
from unittest import mock
from unittest.mock import Mock

def test_sum_geom():
    assert core.sum_geom(1, 1, 1) == 1
    assert core.sum_geom(1, 1, 3) == 3
    assert core.sum_geom(3, 10, 4) == 3333
    assert core.sum_geom(0, 2, 3) == 0
    assert core.sum_geom(1, 0, 3) == 1
    assert core.sum_geom(1, 2, 0) == 0

def test_map_none():
  def fn(x): return x
  assert core.map_none(None, fn) == None
  assert core.map_none("not none", fn) == "not none"

def test_delistify():
  assert core.delistify([1]) == 1
  assert core.delistify((1)) == 1
  assert core.delistify("non list") == "non list"
  assert core.delistify(object) == object

  with pytest.raises(IndexError):
    assert core.delistify([])

def test_datafy():
  x = Mock(data={})
  assert core.datafy(x) == {}
  assert core.datafy([x]) == [{}]
  assert core.datafy([x, x]) == [{}, {}]

@mock.patch("fastai.core.torch.cuda.HalfTensor")
def test_T(HalfTensorMock):
  tensor = torch.ones([1, 2])
  assert core.T(tensor) is tensor

  array = np.arange(0, 5)
  assert core.T(array.astype(np.int)).type() == "torch.LongTensor"
  assert core.T(array.astype(np.float)).type() == "torch.FloatTensor"

  core.T(array.astype(np.float), half=True)
  HalfTensorMock.assert_called_once()

  with pytest.raises(NotImplementedError):
    assert core.T(array.astype(np.object))

def test_create_variable_passing_Variable_object():
  v = torch.autograd.Variable(core.T(np.arange(0, 3)))
  assert core.create_variable(v, volatile=True) is v

@mock.patch("fastai.core.Variable")
def test_create_variable(VariableMock):
  v = np.arange(0, 3)

  with mock.patch("fastai.core.IS_TORCH_04", True):
    core.create_variable(v, volatile=True)
    assert VariableMock.call_args[1] == {"requires_grad": False}

  with mock.patch("fastai.core.IS_TORCH_04", False):
    core.create_variable(v, volatile=True)
    assert VariableMock.call_args[1] == {"requires_grad": False, "volatile": True}

@mock.patch("fastai.core.create_variable")
def test_V_(create_variable_mock):
  core.V_("foo")

  create_variable_mock.assert_called_with('foo', requires_grad=False, volatile=False)

@mock.patch("fastai.core.map_over")
def test_V(map_over_mock):
  core.V("foo")

  assert map_over_mock.call_args[0][0] == 'foo'
  assert type(map_over_mock.call_args[0][1]) == type(lambda:0)

def test_to_np():
  array = np.arange(0, 3).astype(np.float)
  assert core.to_np(array) is array

  tensor = core.T(array)
  result = core.to_np([tensor, tensor])
  np.testing.assert_equal(result[0], array)
  np.testing.assert_equal(result[1], array)

  variable = core.V(array)
  np.testing.assert_equal(core.to_np(variable), array)

def test_noop():
  assert core.noop() is None

def test_partition_functionality():

  def test_partition(a, sz, ex):
    result = core.partition(a, sz)
    assert len(result) == len(ex)
    assert all([a == b for a, b in zip(result, ex)])

  a = [1,2,3,4,5]

  sz = 2
  ex = [[1,2],[3,4],[5]]
  test_partition(a, sz, ex)

  sz = 3
  ex = [[1,2,3],[4,5]]
  test_partition(a, sz, ex)

  sz = 1
  ex = [[1],[2],[3],[4],[5]]
  test_partition(a, sz, ex)

  sz = 6
  ex = [[1,2,3,4,5]]
  test_partition(a, sz, ex)

  sz = 3
  a = []
  result = core.partition(a, sz)
  assert len(result) == 0

def test_partition_error_handling():
  sz = 0
  a = [1,2,3,4,5]
  with pytest.raises(ValueError):
    core.partition(a, sz)

def test_split_by_idxs_functionality():

  seq = [1,2,3,4,5,6]

  def test_split_by_idxs(seq, idxs, ex):
    test_result = []
    for item in core.split_by_idxs(seq, idxs):
      test_result.append(item)

    assert len(test_result) == len(ex)
    assert all([a == b for a,b in zip(test_result, ex)])

  idxs = [2]
  ex = [[1,2],[3,4,5,6]]

  test_split_by_idxs(seq, idxs, ex)

  idxs = [1,2]
  ex = [[1],[2],[3,4,5,6]]
  test_split_by_idxs(seq, idxs, ex)

  idxs = [2,4,5]
  ex = [[1,2],[3,4],[5],[6]]
  test_split_by_idxs(seq, idxs, ex)

  idxs = []
  ex = [[1,2,3,4,5,6]]
  test_split_by_idxs(seq, idxs, ex)

def test_split_by_idxs_error_handling():
  seq = [1,2,3,4]
  idxs = [5]

  gen = core.split_by_idxs(seq, idxs)
  with pytest.raises(KeyError):
    next(gen)


import pandas as pd
import numpy as np

# This case is handled identically to a dict of arrays.
data = np.zeros(shape=(2,), dtype=[('A', 'i4'), ('B', 'f4'), ('C', 'a10')])
data[:] = [(1, 2, 'Hello'), (3., 4., "World")]
df = pd.DataFrame(data=data)
print("DataFrame from Structured array : ")
# we get third column output as byte
print(df)
df = pd.DataFrame(data=data, index=['first', 'second'])
print("DataFrame from Structured array with index : ")
print(df)
pd.DataFrame(data, index=['first', 'second'], columns=['C', 'A', 'B'])
print("DataFrame from Structured array with index and columns : ")
# DataFrame is not intended to work exactly like a 2-dimensional NumPy ndarray.
print(df)


# File di testing per le funzioni di texpy

import texpy as tp
import numpy as np
import unittest

class TestSum(unittest.TestCase):
# Non testo con matrici o vettori le funzioni vettorizzate in quanto si suppone
# che numpy faccia il suo lavoro bene
    def test_notazione_scientifica(self):
        self.assertEqual(tp.ns(123), "$1.23 \\times 10^{2}$")
        self.assertEqual(tp.ns(12.3), "$1.23 \\times 10^{1}$")
        self.assertEqual(tp.ns(382387e-13), "$3.82 \\times 10^{-8}$")
        self.assertEqual(tp.ns(382387e-13, 1e-13), "$382387 \\times 10^{-13}$")
        self.assertEqual(tp.ns(382387e-13, 1e4), "$0.000000000004 \\times 10^{4}$")
        self.assertEqual(tp.ns(123, 1, 1), "$123$")
        self.assertEqual(tp.ns(123, 10, 1), "$12.3 \\times 10^{1}$")
        self.assertEqual(tp.ns(382387e-13, nult=1e-9), "$3.8 \\times 10^{-8}$")
        self.assertEqual(tp.ns(382387e14, nult=1e24), "$4 \\times 10^{19}$")
        self.assertEqual(tp.ns(382387e14, nult=1e4), "$3.823870000000000 \\times 10^{19}$")
        self.assertEqual(tp.ns(382387e14, 1e20), "$0.4 \\times 10^{20}$")
        print("TEST FUNZIONE NOTAZIONE SCIENTIFICA PASSATI")

    def test_numero_errore(self):
        self.assertEqual(tp.ne(1, 0.2), "$1.0 \\pm 0.2$")
        self.assertEqual(tp.ne(1, 20), "$<2 \\times 10^{1}$")
        self.assertEqual(tp.ne(1.987987, 0.2), "$2.0 \\pm 0.2$")
        self.assertEqual(tp.ne(123, 2, unit="F"), "$(123 \\pm 2)$F")
        self.assertEqual(tp.ne(123e-1, 2, "F"), "$(12 \\pm 2)$F")
        self.assertEqual(tp.ne(123e-2, 2, "F"), "$<2$F")
        self.assertEqual(tp.ne(123e-2, 2e-5, "F"), "$(1.23000 \\pm 0.00002)$F")
        self.assertEqual(tp.ne(123e-3, 2e-5, "F"), "$(123.00 \\pm 0.02)$mF")
        #self.assertEqual(tp.ne(0, 0, unit="F"), "$(0 \\pm 0)$F")
        print("\tTEST FUNZIONE NUMERO ERRORE PASSATI")

    def  test_nes(self):
        args = [12e-9, 65, 98e9, 64543]
        output = ["$1.20 \\times 10^{-8}$", "$6.50 \\times 10^{1}$", "$9.80 \\times 10^{10}$", "$6.45 \\times 10^{4}$"]
        for i in range(len(args)):
            self.assertEqual(tp.nes(args[i]), output[i])

        args = [[12.675765e-9, 1e-6], [65.82736, 10], [98.827368e9, 1e1]]
        output = [[0, '$1 \\times 10^{-6}$'],
            ["$7 \\times 10^{1}$", "$1 \\times 10^{1}$"],
            ["$9.882736800 \\times 10^{10}$", "$0.000000001 \\times 10^{10}$"]]
        for i in range(len(args)):
            temp = tp.nes(*args[i])
            for j in range(len(temp)):
                self.assertEqual(temp[j], output[i][j])

        args = [[12.675765e-9, 1e-6, "F"], [65.82736, None, "F"], [98.827368e9, 1e1, "F"]]
        output = [(np.array([0]), np.array(["$1$\\mu F"])),
            np.array(["$66$F"]),
            (np.array(['$98.82736800$GF']), np.array(['$0.00000001$GF']))]
        for i in range(len(args)):
            self.assertEqual(tp.nes(*args[i]), output[i])
        print("\tTEST FUNZIONE NES PASSATI")

    # Per la funzione matrice latex non ho scritto un codice di testing
    # perchè restituisce un output molto lungo

if __name__ == '__main__':
    unittest.main()