# -*- coding: utf-8 -*-
"""Test functions for data input and output."""

__author__ = [
    "SebasKoel",
    "Emiliathewolf",
    "TonyBagnall",
    "jasonlines",
]

__all__ = []

import os
import tempfile

import numpy as np
import pandas as pd
import pytest
from pandas._testing import assert_frame_equal

import sktime
from sktime.datasets import (
    generate_example_long_table,
    load_from_long_to_dataframe,
    load_from_tsfile,
    load_from_tsfile_to_dataframe,
    load_uschange,
    write_dataframe_to_tsfile,
)
from sktime.datasets._data_io import MODULE

def test_load_from_tsfile():
    """Test function for loading TS formats.

    Test
    1. Univariate equal length (UnitTest) returns 2D numpy X, 1D numpy y
    2. Multivariate equal length (BasicMotions) returns 3D numpy X, 1D numpy y
    3. Univariate and multivariate unequal length (PLAID) return X as DataFrame
    """
    data_path = MODULE + "/data/UnitTest/UnitTest_TRAIN.ts"
    # Test 1.1: load univariate equal length (UnitTest), should return 2D array and 1D
    # array, test first and last data
    # Test 1.2: Load a problem without y values (UnitTest),  test first and last data.
    X, y = load_from_tsfile(data_path, return_data_type="np2D")
    X2 = load_from_tsfile(data_path, return_y=False, return_data_type="np2D")
    assert isinstance(X, np.ndarray) and isinstance(y, np.ndarray)
    assert X.ndim == 2 and X2.ndim == 2
    assert X.shape == (20, 24) and y.shape == (20,)
    assert X[0][0] == 573.0
    X2 = load_from_tsfile(data_path, return_y=False, return_data_type="numpy3D")
    assert isinstance(X2, np.ndarray)
    assert X2.ndim == 3
    assert X2.shape == (20, 1, 24)
    assert X2[0][0][0] == 573.0

    # Test 2: load multivare equal length (BasicMotions), should return 3D array and 1D
    # array, test first and last data.
    data_path = MODULE + "/data/BasicMotions/BasicMotions_TRAIN.ts"
    X, y = load_from_tsfile(data_path, return_data_type="numpy3d")
    assert isinstance(X, np.ndarray) and isinstance(y, np.ndarray)
    assert X.shape == (40, 6, 100) and y.shape == (40,)
    assert X[1][2][3] == -1.898794
    X, y = load_from_tsfile(data_path)
    assert isinstance(X, pd.DataFrame) and isinstance(y, np.ndarray)
    assert X.shape == (40, 6) and y.shape == (40,)
    assert isinstance(X.iloc[1, 2], pd.Series)
    assert X.iloc[1, 2].iloc[3] == -1.898794

    # Test 3.1: load univariate unequal length (PLAID), should return a one column
    # dataframe,
    data_path = MODULE + "/data/PLAID/PLAID_TRAIN.ts"
    X, y = load_from_tsfile(full_file_path_and_name=data_path)
    assert isinstance(X, pd.DataFrame) and isinstance(y, np.ndarray)
    assert X.shape == (537, 1) and y.shape == (537,)
    # Test 3.2: load multivariate unequal length (JapaneseVowels), should return a X
    # columns dataframe,
    data_path = MODULE + "/data/JapaneseVowels/JapaneseVowels_TRAIN.ts"
    X, y = load_from_tsfile(full_file_path_and_name=data_path)
    assert isinstance(X, pd.DataFrame) and isinstance(y, np.ndarray)
    assert X.shape == (270, 12) and y.shape == (270,)

_CHECKS = {
    "uschange": {
        "columns": ["Income", "Production", "Savings", "Unemployment"],
        "len_y": 187,
        "len_X": 187,
        "data_types_X": {
            "Income": "float64",
            "Production": "float64",
            "Savings": "float64",
            "Unemployment": "float64",
        },
        "data_type_y": "float64",
        "data": load_uschange(),
    },
}

@pytest.mark.parametrize("dataset", sorted(_CHECKS.keys()))
def test_data_loaders(dataset):
    """
    Assert if datasets are loaded correctly.

    dataset: dictionary with values to assert against should contain:
        'columns' : list with column names in correct order,
        'len_y'   : lenght of the y series (int),
        'len_X'   : lenght of the X series/dataframe (int),
        'data_types_X' : dictionary with column name keys and dtype as value,
        'data_type_y'  : dtype if y column (string)
        'data'    : tuple with y series and X series/dataframe if one is not
                    applicable fill with None value,
    """
    checks = _CHECKS[dataset]
    y = checks["data"][0]
    X = checks["data"][1]

    if y is not None:
        assert isinstance(y, pd.Series)
        assert len(y) == checks["len_y"]
        assert y.dtype == checks["data_type_y"]

    if X is not None:
        if len(checks["data_types_X"]) > 1:
            assert isinstance(X, pd.DataFrame)
        else:
            assert isinstance(X, pd.Series)

        assert X.columns.values.tolist() == checks["columns"]

        for col, dt in checks["data_types_X"].items():
            assert X[col].dtype == dt

        assert len(X) == checks["len_X"]

def test_load_from_tsfile_to_dataframe():
    """Test the load_from_tsfile_to_dataframe() function."""
    # Test that an empty file is classed an invalid
    fd, path = tempfile.mkstemp()
    try:
        with os.fdopen(fd, "w") as tmp_file:
            # Write the contents of the file
            file_contents = ""
            tmp_file.write(file_contents)
            tmp_file.flush()
            # Parse the file and assert that it is invalid
            np.testing.assert_raises(IOError, load_from_tsfile_to_dataframe, path)
    finally:
        os.remove(path)
    # Test that a file with an incomplete set of metadata is invalid
    fd, path = tempfile.mkstemp()
    try:
        with os.fdopen(fd, "w") as tmp_file:
            # Write the contents of the file
            file_contents = (
                "@problemName Test Problem\n@timeStamps " "true\n@univariate true\n"
            )
            tmp_file.write(file_contents)
            tmp_file.flush()
            # Parse the file and assert that it is invalid
            np.testing.assert_raises(IOError, load_from_tsfile_to_dataframe, path)
    finally:
        os.remove(path)
    # Test that a file with a complete set of metadata but no data is invalid
    fd, path = tempfile.mkstemp()
    try:
        with os.fdopen(fd, "w") as tmp_file:
            # Write the contents of the file
            file_contents = (
                "@problemName Test Problem\n@timeStamps "
                "true\n@univariate true\n@classLabel false\n@data"
            )
            tmp_file.write(file_contents)
            tmp_file.flush()
            # Parse the file and assert that it is invalid
            np.testing.assert_raises(IOError, load_from_tsfile_to_dataframe, path)
    finally:
        os.remove(path)
    # Test that a file with a complete set of metadata and no data but
    # invalid metadata values is invalid
    fd, path = tempfile.mkstemp()
    try:
        with os.fdopen(fd, "w") as tmp_file:
            # Write the contents of the file
            file_contents = (
                "@problemName\n@timeStamps\n@univariate "
                "true\n@classLabel false\n@data"
            )
            tmp_file.write(file_contents)
            tmp_file.flush()
            # Parse the file and assert that it is invalid
            np.testing.assert_raises(IOError, load_from_tsfile_to_dataframe, path)
    finally:
        os.remove(path)
    # Test that a file with a complete set of metadata and a single
    # case/dimension parses correctly
    fd, path = tempfile.mkstemp()
    try:
        with os.fdopen(fd, "w") as tmp_file:
            # Write the contents of the file
            file_contents = (
                "@problemName Test Problem\n@timeStamps "
                "true\n@univariate true\n@classLabel "
                "false\n@data\n"
            )
            file_contents += "(0, 1), (1, 2)"

            tmp_file.write(file_contents)
            tmp_file.flush()
            # Parse the file
            df = load_from_tsfile_to_dataframe(path)
            # Test the DataFrame returned accurately reflects the data in
            # the file
            np.testing.assert_equal(len(df), 1)
            np.testing.assert_equal(len(df.columns), 1)
            series = df["dim_0"]
            np.testing.assert_equal(len(series), 1)
            series = df["dim_0"][0]
            np.testing.assert_equal(series[0], 1.0)
            np.testing.assert_equal(series[1], 2.0)
    finally:
        os.remove(path)
    # Test that a file with a complete set of metadata and 2 cases with 3
    # dimensions parses correctly
    fd, path = tempfile.mkstemp()
    try:
        with os.fdopen(fd, "w") as tmp_file:
            # Write the contents of the file
            file_contents = (
                "@problemName Test Problem\n@timeStamps "
                "true\n@univariate true\n@classLabel "
                "false\n@data\n"
            )
            file_contents += "(0, 1), (1, 2):(0, 3), (1, 4):(0, 5), (1, 6)\n"
            file_contents += "(0, 11), (1, 12):(0, 13), (1,14):(0, 15), (1, 16)     \n"
            tmp_file.write(file_contents)
            tmp_file.flush()
            # Parse the file
            df = load_from_tsfile_to_dataframe(path)
            # Test the DataFrame returned accurately reflects the data in
            # the file
            np.testing.assert_equal(len(df), 2)
            np.testing.assert_equal(len(df.columns), 3)
            series = df["dim_0"]
            np.testing.assert_equal(len(series), 2)

            series = df["dim_0"][0]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 1.0)
            np.testing.assert_equal(series[1], 2.0)

            series = df["dim_0"][1]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 11.0)
            np.testing.assert_equal(series[1], 12.0)

            series = df["dim_1"]
            np.testing.assert_equal(len(series), 2)

            series = df["dim_1"][0]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 3.0)
            np.testing.assert_equal(series[1], 4.0)

            series = df["dim_1"][1]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 13.0)
            np.testing.assert_equal(series[1], 14.0)

            series = df["dim_2"]
            np.testing.assert_equal(len(series), 2)

            series = df["dim_2"][0]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 5.0)
            np.testing.assert_equal(series[1], 6.0)

            series = df["dim_2"][1]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 15.0)
            np.testing.assert_equal(series[1], 16.0)
    finally:
        os.remove(path)
    # Test that a file with a complete set of metadata and time-series of
    # different length parses correctly
    fd, path = tempfile.mkstemp()
    try:
        with os.fdopen(fd, "w") as tmp_file:
            # Write the contents of the file
            file_contents = (
                "@problemName Test Problem\n@timeStamps "
                "true\n@univariate true\n@classLabel "
                "false\n@data\n"
            )
            file_contents += "(0, 1), (1, 2):(0, 3):(0, 5), (1, 6)\n"
            file_contents += "(0, 11), (1, 12):(0, 13), (1,14):(0, 15)\n"
            tmp_file.write(file_contents)
            tmp_file.flush()
            # Parse the file
            df = load_from_tsfile_to_dataframe(path)
            # Test the DataFrame returned accurately reflects the data in
            # the file

            np.testing.assert_equal(len(df), 2)
            np.testing.assert_equal(len(df.columns), 3)

            series = df["dim_0"]
            np.testing.assert_equal(len(series), 2)

            series = df["dim_0"][0]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 1.0)
            np.testing.assert_equal(series[1], 2.0)

            series = df["dim_0"][1]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 11.0)
            np.testing.assert_equal(series[1], 12.0)

            series = df["dim_1"]
            np.testing.assert_equal(len(series), 2)

            series = df["dim_1"][0]
            np.testing.assert_equal(len(series), 1)
            np.testing.assert_equal(series[0], 3.0)

            series = df["dim_1"][1]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 13.0)
            np.testing.assert_equal(series[1], 14.0)

            series = df["dim_2"]
            np.testing.assert_equal(len(series), 2)

            series = df["dim_2"][0]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 5.0)
            np.testing.assert_equal(series[1], 6.0)

            series = df["dim_2"][1]
            np.testing.assert_equal(len(series), 1)
            np.testing.assert_equal(series[0], 15.0)

    finally:
        os.remove(path)

    # Test that a file with a complete set of metadata and data but an
    # inconsistent number of dimensions across cases is classed as invalid

    fd, path = tempfile.mkstemp()
    try:
        with os.fdopen(fd, "w") as tmp_file:

            # Write the contents of the file

            file_contents = (
                "@problemName Test Problem\n@timeStamps "
                "true\n@univariate true\n@classLabel "
                "false\n@data\n"
            )
            file_contents += "(0, 1), (1, 2):(0, 3), (1, 4):(0, 5), (1, 6)\n"
            file_contents += "(0, 11), (1, 12):(0, 13), (1,14)    \n"

            tmp_file.write(file_contents)
            tmp_file.flush()

            # Parse the file and assert that it is invalid

            np.testing.assert_raises(IOError, load_from_tsfile_to_dataframe, path)

    finally:
        os.remove(path)

    # Test that a file with a complete set of metadata and data but missing
    # values after a tuple is classed as invalid

    fd, path = tempfile.mkstemp()
    try:
        with os.fdopen(fd, "w") as tmp_file:

            # Write the contents of the file

            file_contents = (
                "@problemName Test Problem\n@timeStamps "
                "true\n@univariate true\n@classLabel "
                "false\n@data\n"
            )
            file_contents += "(0, 1), (1, 2):(0, 3), (1, 4):(0, 5),\n"

            tmp_file.write(file_contents)
            tmp_file.flush()

            # Parse the file and assert that it is invalid

            np.testing.assert_raises(IOError, load_from_tsfile_to_dataframe, path)

    finally:
        os.remove(path)

    # Test that a file with a complete set of metadata and data and some
    # empty dimensions is classed as valid

    fd, path = tempfile.mkstemp()
    try:
        with os.fdopen(fd, "w") as tmp_file:

            # Write the contents of the file

            file_contents = (
                "@problemName Test Problem\n@timeStamps "
                "true\n@univariate true\n@classLabel "
                "false\n@data\n"
            )
            file_contents += "(0, 1), (1, 2):     :(0, 5), (1, 6)\n"
            file_contents += "(0, 11), (1, 12):(0, 13), (1,14)    :       \n"
            file_contents += (
                "(0, 21), (1, 22):(0, 23), (1,24)    :   (0,25), (1, 26)    \n"
            )

            tmp_file.write(file_contents)
            tmp_file.flush()

            # Parse the file

            df = load_from_tsfile_to_dataframe(path)

            # Test the DataFrame returned accurately reflects the data in
            # the file

            np.testing.assert_equal(len(df), 3)
            np.testing.assert_equal(len(df.columns), 3)

            series = df["dim_0"]
            np.testing.assert_equal(len(series), 3)

            series = df["dim_0"][0]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 1.0)
            np.testing.assert_equal(series[1], 2.0)

            series = df["dim_0"][1]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 11.0)
            np.testing.assert_equal(series[1], 12.0)

            series = df["dim_0"][2]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 21.0)
            np.testing.assert_equal(series[1], 22.0)

            series = df["dim_1"]
            np.testing.assert_equal(len(series), 3)

            series = df["dim_1"][0]
            np.testing.assert_equal(len(series), 0)

            series = df["dim_1"][1]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 13.0)
            np.testing.assert_equal(series[1], 14.0)

            series = df["dim_1"][2]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 23.0)
            np.testing.assert_equal(series[1], 24.0)

            series = df["dim_2"]
            np.testing.assert_equal(len(series), 3)

            series = df["dim_2"][0]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 5.0)
            np.testing.assert_equal(series[1], 6.0)

            series = df["dim_2"][1]
            np.testing.assert_equal(len(series), 0)

            series = df["dim_2"][2]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 25.0)
            np.testing.assert_equal(series[1], 26.0)

    finally:
        os.remove(path)

    # Test that a file with a complete set of metadata and data that
    # contains datetimes as timestamps and has some empty dimensions is
    # classed as valid

    fd, path = tempfile.mkstemp()
    try:
        with os.fdopen(fd, "w") as tmp_file:

            # Write the contents of the file

            file_contents = (
                "@problemName Test Problem\n@timeStamps "
                "true\n@univariate true\n@classLabel "
                "false\n@data\n"
            )
            file_contents += (
                "(01/01/2019 00:00:00, 1),  (01/02/2019 "
                "00:00:00, 2)  :                               "
                "                      : (01/05/2019 00:00:00, "
                "5), (01/06/2019 00:00:00, 6)\n"
            )
            file_contents += (
                "(01/01/2020 00:00:00, 11), (01/02/2020 "
                "00:00:00, 12) : (01/03/2020 00:00:00, 13), "
                "(01/04/2020 00:00:00, 14) :  \n"
            )
            file_contents += (
                "(01/01/2021 00:00:00, 21), (01/02/2021 "
                "00:00:00, 22) : (01/03/2021 00:00:00, 23), "
                "(01/04/2021 00:00:00, 24) :  \n"
            )

            tmp_file.write(file_contents)
            tmp_file.flush()

            # Parse the file

            df = load_from_tsfile_to_dataframe(path)

            # Test the DataFrame returned accurately reflects the data in
            # the file

            np.testing.assert_equal(len(df), 3)
            np.testing.assert_equal(len(df.columns), 3)

            series = df["dim_0"]
            np.testing.assert_equal(len(series), 3)

            series = df["dim_0"][0]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series["01/01/2019"], 1.0)
            np.testing.assert_equal(series["01/02/2019"], 2.0)

            series = df["dim_0"][1]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series["01/01/2020"], 11.0)
            np.testing.assert_equal(series["01/02/2020"], 12.0)

            series = df["dim_0"][2]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series["01/01/2021"], 21.0)
            np.testing.assert_equal(series["01/02/2021"], 22.0)

            series = df["dim_1"]
            np.testing.assert_equal(len(series), 3)

            series = df["dim_1"][0]
            np.testing.assert_equal(len(series), 0)

            series = df["dim_1"][1]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series["01/03/2020"], 13.0)
            np.testing.assert_equal(series["01/04/2020"], 14.0)

            series = df["dim_1"][2]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series["01/03/2021"], 23.0)
            np.testing.assert_equal(series["01/04/2021"], 24.0)

            series = df["dim_2"]
            np.testing.assert_equal(len(series), 3)

            series = df["dim_2"][0]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series["01/05/2019"], 5.0)
            np.testing.assert_equal(series["01/06/2019"], 6.0)

            series = df["dim_2"][1]
            np.testing.assert_equal(len(series), 0)

            series = df["dim_2"][2]
            np.testing.assert_equal(len(series), 0)

    finally:
        os.remove(path)

    # Test that a file that mixes timestamp conventions is invalid

    fd, path = tempfile.mkstemp()
    try:
        with os.fdopen(fd, "w") as tmp_file:

            # Write the contents of the file

            file_contents = (
                "@problemName Test Problem\n@timeStamps "
                "true\n@univariate true\n@classLabel "
                "false\n@data\n"
            )
            file_contents += (
                "(01/01/2019 00:00:00, 1),  (01/02/2019 "
                "00:00:00, 2)  :                               "
                "                      : (01/05/2019 00:00:00, "
                "5), (01/06/2019 00:00:00, 6)\n"
            )
            file_contents += (
                "(00, 11), (1, 12) : (01/03/2020 00:00:00, 13), "
                "(01/04/2020 00:00:00, 14) :  \n"
            )
            file_contents += (
                "(01/01/2021 00:00:00, 21), (01/02/2021 "
                "00:00:00, 22) : (01/03/2021 00:00:00, 23), "
                "(01/04/2021 00:00:00, 24) :  \n"
            )

            tmp_file.write(file_contents)
            tmp_file.flush()

            # Parse the file and assert that it is invalid

            np.testing.assert_raises(IOError, load_from_tsfile_to_dataframe, path)

    finally:
        os.remove(path)

    # Test that a file with a complete set of metadata and data but missing
    # classes is classed as invalid

    fd, path = tempfile.mkstemp()
    try:
        with os.fdopen(fd, "w") as tmp_file:

            # Write the contents of the file

            file_contents = (
                "@problemName Test Problem\n@timeStamps "
                "true\n@univariate true\n@classLabel true 0 1 "
                "2\n@data\n"
            )
            file_contents += "(0, 1), (1, 2):(0, 3), (1, 4):(0, 5), (1, 6)\n"
            file_contents += "(0, 11), (1, 12):(0, 13), (1,14):(0, 15), (1, 16)     \n"

            tmp_file.write(file_contents)
            tmp_file.flush()

            # Parse the file and assert that it is invalid

            np.testing.assert_raises(IOError, load_from_tsfile_to_dataframe, path)

    finally:
        os.remove(path)

    # Test that a file with a complete set of metadata and data but invalid
    # classes is classed as invalid

    fd, path = tempfile.mkstemp()
    try:
        with os.fdopen(fd, "w") as tmp_file:

            # Write the contents of the file

            file_contents = (
                "@problemName Test Problem\n@timeStamps "
                "true\n@univariate true\n@classLabel true 0 1 "
                "2\n@data\n"
            )
            file_contents += "(0, 1), (1, 2):(0, 3), (1, 4):(0, 5), (1, 6) : 0 \n"
            file_contents += (
                "(0, 11), (1, 12):(0, 13), (1,14):(0, 15), (1, 16)   : 3  \n"
            )

            tmp_file.write(file_contents)
            tmp_file.flush()

            # Parse the file and assert that it is invalid

            np.testing.assert_raises(IOError, load_from_tsfile_to_dataframe, path)

    finally:
        os.remove(path)

    # Test that a file with a complete set of metadata and data with classes
    # is classed as valid

    fd, path = tempfile.mkstemp()
    try:
        with os.fdopen(fd, "w") as tmp_file:

            # Write the contents of the file

            file_contents = (
                "@problemName Test Problem\n@timeStamps "
                "true\n@univariate true\n@classLabel true 0 1 "
                "2\n@data\n"
            )
            file_contents += "(0, 1), (1, 2):(0, 3), (1, 4):(0, 5), (1, 6): 0\n"
            file_contents += (
                "(0, 11), (1, 12):(0, 13), (1,14):(0, 15), (1, 16): 2     \n"
            )

            tmp_file.write(file_contents)
            tmp_file.flush()

            # Parse the file

            df, y = load_from_tsfile_to_dataframe(path)

            # Test the DataFrame of X values returned accurately reflects
            # the data in the file

            np.testing.assert_equal(len(df), 2)
            np.testing.assert_equal(len(df.columns), 3)

            series = df["dim_0"]
            np.testing.assert_equal(len(series), 2)

            series = df["dim_0"][0]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 1.0)
            np.testing.assert_equal(series[1], 2.0)

            series = df["dim_0"][1]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 11.0)
            np.testing.assert_equal(series[1], 12.0)

            series = df["dim_1"]
            np.testing.assert_equal(len(series), 2)

            series = df["dim_1"][0]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 3.0)
            np.testing.assert_equal(series[1], 4.0)

            series = df["dim_1"][1]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 13.0)
            np.testing.assert_equal(series[1], 14.0)

            series = df["dim_2"]
            np.testing.assert_equal(len(series), 2)

            series = df["dim_2"][0]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 5.0)
            np.testing.assert_equal(series[1], 6.0)

            series = df["dim_2"][1]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 15.0)
            np.testing.assert_equal(series[1], 16.0)

            # Test that the class values are as expected

            np.testing.assert_equal(len(y), 2)
            np.testing.assert_equal(y[0], "0")
            np.testing.assert_equal(y[1], "2")

    finally:
        os.remove(path)

    # Test that a file with a complete set of metadata and data, with no
    # timestamps, is classed as valid

    fd, path = tempfile.mkstemp()
    try:
        with os.fdopen(fd, "w") as tmp_file:

            # Write the contents of the file

            file_contents = (
                "@problemName Test Problem\n@timeStamps "
                "false\n@univariate true\n@classLabel "
                "false\n@data\n"
            )
            file_contents += "1,2:3,4:5,6\n"
            file_contents += "11,12:13,14:15,16\n"
            file_contents += "21,22:23,24:25,26\n"

            tmp_file.write(file_contents)
            tmp_file.flush()

            # Parse the file

            df = load_from_tsfile_to_dataframe(path)

            # Test the DataFrame returned accurately reflects the data in
            # the file

            np.testing.assert_equal(len(df), 3)
            np.testing.assert_equal(len(df.columns), 3)

            series = df["dim_0"]
            np.testing.assert_equal(len(series), 3)

            series = df["dim_0"][0]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 1.0)
            np.testing.assert_equal(series[1], 2.0)

            series = df["dim_0"][1]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 11.0)
            np.testing.assert_equal(series[1], 12.0)

            series = df["dim_0"][2]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 21.0)
            np.testing.assert_equal(series[1], 22.0)

            series = df["dim_1"]
            np.testing.assert_equal(len(series), 3)

            series = df["dim_1"][0]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 3.0)
            np.testing.assert_equal(series[1], 4.0)

            series = df["dim_1"][1]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 13.0)
            np.testing.assert_equal(series[1], 14.0)

            series = df["dim_1"][2]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 23.0)
            np.testing.assert_equal(series[1], 24.0)

            series = df["dim_2"]
            np.testing.assert_equal(len(series), 3)

            series = df["dim_2"][0]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 5.0)
            np.testing.assert_equal(series[1], 6.0)

            series = df["dim_2"][1]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 15.0)
            np.testing.assert_equal(series[1], 16.0)

            series = df["dim_2"][2]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 25.0)
            np.testing.assert_equal(series[1], 26.0)

    finally:
        os.remove(path)

    # Test that a file with a complete set of metadata and data, with no
    # timestamps and some empty dimensions, is classed as valid

    fd, path = tempfile.mkstemp()
    try:
        with os.fdopen(fd, "w") as tmp_file:

            # Write the contents of the file

            file_contents = (
                "@problemName Test Problem\n@timeStamps "
                "false\n@univariate true\n@classLabel "
                "false\n@data\n"
            )
            file_contents += "1,2::5,6\n"
            file_contents += "11,12:13,14:15,16\n"
            file_contents += "21,22:23,24:\n"

            tmp_file.write(file_contents)
            tmp_file.flush()

            # Parse the file

            df = load_from_tsfile_to_dataframe(path)

            # Test the DataFrame returned accurately reflects the data in
            # the file

            np.testing.assert_equal(len(df), 3)
            np.testing.assert_equal(len(df.columns), 3)

            series = df["dim_0"]
            np.testing.assert_equal(len(series), 3)

            series = df["dim_0"][0]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 1.0)
            np.testing.assert_equal(series[1], 2.0)

            series = df["dim_0"][1]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 11.0)
            np.testing.assert_equal(series[1], 12.0)

            series = df["dim_0"][2]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 21.0)
            np.testing.assert_equal(series[1], 22.0)

            series = df["dim_1"]
            np.testing.assert_equal(len(series), 3)

            series = df["dim_1"][0]
            np.testing.assert_equal(len(series), 0)

            series = df["dim_1"][1]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 13.0)
            np.testing.assert_equal(series[1], 14.0)

            series = df["dim_1"][2]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 23.0)
            np.testing.assert_equal(series[1], 24.0)

            series = df["dim_2"]
            np.testing.assert_equal(len(series), 3)

            series = df["dim_2"][0]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 5.0)
            np.testing.assert_equal(series[1], 6.0)

            series = df["dim_2"][1]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 15.0)
            np.testing.assert_equal(series[1], 16.0)

            series = df["dim_2"][2]
            np.testing.assert_equal(len(series), 0)

    finally:
        os.remove(path)

    # Test that a file with a complete set of metadata and data, with no
    # timestamps and some empty dimensions and classes, is classed as valid

    fd, path = tempfile.mkstemp()
    try:
        with os.fdopen(fd, "w") as tmp_file:
            # Write the contents of the file

            file_contents = (
                "@problemName Test Problem\n@timeStamps "
                "false\n@univariate true\n@classLabel true cat "
                "bear dog\n@data\n"
            )
            file_contents += "1,2::5,6:cat  \n"
            file_contents += "11,12:13,14:15,16:  dog\n"
            file_contents += "21,22:23,24::   bear   \n"

            tmp_file.write(file_contents)
            tmp_file.flush()

            # Parse the file

            df, y = load_from_tsfile_to_dataframe(path)

            # Test the DataFrame of X values returned accurately reflects
            # the data in the file

            np.testing.assert_equal(len(df), 3)
            np.testing.assert_equal(len(df.columns), 3)

            series = df["dim_0"]
            np.testing.assert_equal(len(series), 3)

            series = df["dim_0"][0]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 1.0)
            np.testing.assert_equal(series[1], 2.0)

            series = df["dim_0"][1]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 11.0)
            np.testing.assert_equal(series[1], 12.0)

            series = df["dim_0"][2]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 21.0)
            np.testing.assert_equal(series[1], 22.0)

            series = df["dim_1"]
            np.testing.assert_equal(len(series), 3)

            series = df["dim_1"][0]
            np.testing.assert_equal(len(series), 0)

            series = df["dim_1"][1]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 13.0)
            np.testing.assert_equal(series[1], 14.0)

            series = df["dim_1"][2]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 23.0)
            np.testing.assert_equal(series[1], 24.0)

            series = df["dim_2"]
            np.testing.assert_equal(len(series), 3)

            series = df["dim_2"][0]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 5.0)
            np.testing.assert_equal(series[1], 6.0)

            series = df["dim_2"][1]
            np.testing.assert_equal(len(series), 2)
            np.testing.assert_equal(series[0], 15.0)
            np.testing.assert_equal(series[1], 16.0)

            series = df["dim_2"][2]
            np.testing.assert_equal(len(series), 0)

            # Test that the class values are as expected

            np.testing.assert_equal(len(y), 3)
            np.testing.assert_equal(y[0], "cat")
            np.testing.assert_equal(y[1], "dog")
            np.testing.assert_equal(y[2], "bear")

    finally:
        os.remove(path)

def test_load_from_long_to_dataframe(tmpdir):
    """Test for loading from long to dataframe."""
    # create and save a example long-format file to csv
    test_dataframe = generate_example_long_table()
    dataframe_path = tmpdir.join("data.csv")
    test_dataframe.to_csv(dataframe_path, index=False)
    # load and convert the csv to sktime-formatted data
    nested_dataframe = load_from_long_to_dataframe(dataframe_path)
    assert isinstance(nested_dataframe, pd.DataFrame)

def test_load_from_long_incorrect_format(tmpdir):
    """Test for loading from long with incorrect format."""
    with pytest.raises(ValueError):
        dataframe = generate_example_long_table()
        dataframe.drop(dataframe.columns[[3]], axis=1, inplace=True)
        dataframe_path = tmpdir.join("data.csv")
        dataframe.to_csv(dataframe_path, index=False)
        load_from_long_to_dataframe(dataframe_path)

@pytest.mark.parametrize("dataset", ["ItalyPowerDemand", "BasicMotions"])
def test_write_dataframe_to_ts_success(tmp_path, dataset):
    """Tests whether a dataset can be written by the .ts writer then read in."""
    # load an example dataset
    path = os.path.join(
        os.path.dirname(sktime.__file__),
        f"datasets/data/{dataset}/{dataset}_TEST.ts",
    )
    test_X, test_y = load_from_tsfile_to_dataframe(path)
    # output the dataframe in a ts file
    write_dataframe_to_tsfile(
        data=test_X,
        path=tmp_path,
        problem_name=dataset,
        class_label=np.unique(test_y),
        class_value_list=test_y,
        comment="""
          The data was derived from twelve monthly electrical power demand
          time series from Italy and first used in the paper "Intelligent
          Icons: Integrating Lite-Weight Data Mining and Visualization into
          GUI Operating Systems". The classification task is to distinguish
          days from Oct to March (inclusive) from April to September.
        """,
        fold="_transform",
    )
    # load data back from the ts file
    result = f"{tmp_path}/{dataset}/{dataset}_transform.ts"
    res_X, res_y = load_from_tsfile_to_dataframe(result)
    # check if the dataframes are the same
    assert_frame_equal(res_X, test_X)

def test_write_dataframe_to_ts_fail(tmp_path):
    """Tests if non-dataframes are handled correctly."""
    with pytest.raises(ValueError, match="Data provided must be a DataFrame"):
        write_dataframe_to_tsfile(
            data=np.random.rand(3, 2),
            path=str(tmp_path),
            problem_name="GunPoint",
        )


##### Import packages
# Basic packages
import pandas as pd
import numpy as np

# Data Visualization packages
import matplotlib.pyplot as plt
import seaborn as sns
import folium
from folium import FeatureGroup, LayerControl, Map

# Other packages
import time
from datetime import datetime
import json
import geopandas as gpd

# To avoid warnings
import warnings
warnings.filterwarnings("ignore")

##### Functions

# Create labels on charts
def autolabel(plot):
    for p in plot.patches:
        plot.annotate(format(p.get_height(), '.0f'),
                       (p.get_x() + p.get_width() / 2., p.get_height()),
                       ha = 'center', va = 'center',
                       xytext = (0, 9),
                       textcoords = 'offset points')
# Plot a bar chart
def graph_bar(data,title):
    plt.figure(figsize=(35, 10))
    plot = sns.barplot(x = data.index, y = data.values, palette="rocket")
    plt.title(title)
    plt.xticks(rotation=90)
    autolabel(plot)
    plt.tight_layout()
    plt.show()

# Plot a pie chart
def graph_pie(dictionary,title):
    dictionary = dict(sorted(dictionary.items(), key = lambda item: item[1], reverse=True))
    plt.figure(figsize=(10, 10))
    plt.pie(dictionary.values(), labels = dictionary.keys(), explode = [0.1 for i in range(len(dictionary.values()))],
            autopct='%.1f%%', shadow = True, labeldistance = 1.07, startangle = 45)
    plt.title(title)
    centre_circle = plt.Circle((0,0),0.80,fc='white')
    fig = plt.gcf()
    fig.gca().add_artist(centre_circle)
    plt.axis('equal')
    plt.tight_layout()
    plt.show()

# # Initial analysis

##### Import data
# Check the csv's path before running it

df_est = pd.read_csv("CoordEstados.csv", encoding = "ISO-8859-1") # Mexican states data
df_cov = pd.read_csv("14.11.20 - COVID19MEXICO.csv", encoding = "ISO-8859-1") # Covid-19 data

##### Brief statistic of each DataFrame

# Mexican states data
print(' Mexican states data '.center(35,'#'))
df_est.info()
df_est.describe()
print(' ')

# Covid-19 data
print(' Covid-19 data '.center(35,'#'))
df_cov.info()
df_cov.describe()

##### Number of missing values for each DataFrame

for dataframe,name in zip([df_est,df_cov],['Mexican states data','Covid-19 data']):
    print(f' {name}  nº of NaNs: '.center(55,'#'))
    if dataframe.isna().sum().sum() == 0:
        print('0 NaNs in total dataset')
    else:
        for column in dataframe.columns:
            if dataframe[column].isna().sum() != 0:
                print(f'{column} has {round(dataframe[column].isna().sum()/dataframe.shape[0],3)}% of NaNs')

# # Exploratory Data Analysis

##### Transform all the mexican states from number to name

states_num = sorted(df_cov['ENTIDAD_RES'].unique())
states_name = df_est['Estado'].loc[:31].values

dict_states = dict()
for num,name in zip(states_num,states_name):
    dict_states[num] = name
print(dict_states)

df_cov['ENTIDAD_RES'] = df_cov['ENTIDAD_RES'].map(dict_states)

##### Transform all the mexican National Health System institutions from number to name

dict_sector = {1:'CRUZ ROJA', 2:'DIF',3:'ESTATAL',4:'IMSS',5:'IMSS-BIENESTAR',6:'ISSSTE',7:'MUNICIPAL',8:'PEMEX',
               9:'PRIVADA',10:'SEDENA',11:'SEMAR',12:'SSA',13:'UNIVERSITARIO',99:'NO ESPECIFICADO'}

df_cov['SECTOR'] = df_cov['SECTOR'].map(dict_sector)

##### Create a new colum with the time difference between been positive in COVID-19 and die
# If the person didn't die, time difference is 0

df_cov['FECHA_SINTOMAS'] = df_cov['FECHA_SINTOMAS'].apply(lambda x: datetime.strptime(x, '%Y-%m-%d'))
df_cov['FECHA_DEF'] = df_cov['FECHA_DEF'].replace('9999-99-99', '2001-01-01')

df_cov['FECHA_DEF'] = df_cov['FECHA_DEF'].apply(lambda x: datetime.strptime(x, '%Y-%m-%d'))
df_cov['DIFERENCIA'] = df_cov['FECHA_DEF'].sub(df_cov['FECHA_SINTOMAS'], axis=0)

df_cov['DIFERENCIA'] = df_cov['DIFERENCIA'] / np.timedelta64(1, 'D')
df_cov.loc[df_cov['DIFERENCIA']<0,'DIFERENCIA'] = 0

# ## Positive COVID-19 cases

##### Analysis of total number of positive cases
# Catalogo and Descriptores show all the information about the variables.
# CLASIFICACION_FINAL has 1, 2 or 3 when a person is COVID-19 confirmed.

df_positive_cases =  df_cov[df_cov['CLASIFICACION_FINAL'].isin([1,2,3])]

fig, axes = plt.subplots(1, 4, figsize=(25, 5))

# 1st graph - Gender of positive cases
axes[0].set_title('Gender of positive cases')
plot = sns.barplot(x = ['Women','Men'],y = df_positive_cases['SEXO'].value_counts().values, palette = 'pastel', ax = axes[0])
plot.set(ylim = (450000,550000), ylabel = None, yticklabels = [])
plot.tick_params(left=False)
autolabel(plot)

# 2nd graph - Result of diagnostic (positive or not)
axes[1].set_title('Result of diagnostic')
pos, neg = [df_positive_cases.shape[0], df_cov.query('CLASIFICACION_FINAL == 7').shape[0]]
plot = sns.barplot(x = ['Positive','Negatuve'],y = [pos,neg], palette = 'pastel', ax = axes[1])
plot.set(ylim = (800000,1500000), ylabel = None, yticklabels = [])
plot.tick_params(left=False)
autolabel(plot)

# 3rd graph - Age distribution of positives cases
axes[2].set_title('Age distribution of positive cases')
plot = sns.distplot(df_positive_cases.EDAD, ax = axes[2])
plot.set(ylabel = None, yticklabels = [])
plot.tick_params(left=False)

# 4th graph - National Health System institution that provided the care
axes[3].set_title('National Health System institution \n that provided the care')
data = df_positive_cases.SECTOR.value_counts()
plot = sns.barplot(x = data.index[:5], y = data.values[:5],
                   palette = 'pastel', ax = axes[3])
plot.set(ylim = (10000,650000), ylabel = None, yticklabels = [])
plot.tick_params(left=False)
autolabel(plot)

plt.tight_layout()
plt.show()

##### Total number of positive cases by state

# Data Extraction
print(f' Total number of positive cases = {df_positive_cases.shape[0]} | % of total = {(df_positive_cases.shape[0]/df_cov.shape[0]):.3}% '.center(100,'#'))
data = df_positive_cases.groupby('ENTIDAD_RES').count()['CLASIFICACION_FINAL'].sort_values(ascending = False)

# Data Visualization
graph_bar(data, 'Total number of positive cases by state')

##### Mean age of positive cases by state

# Data Extraction
print(f' Total mean age of positive cases = {(df_positive_cases.EDAD.mean()):.3} '.center(100,'#'))
data = df_positive_cases.groupby('ENTIDAD_RES').mean()['EDAD'].sort_values(ascending = False)

# Data Visualization
graph_bar(data, 'Mean age of positive cases by state')

##### Common illness on total positive cases

# Data Extraction
ill_name = ['DIABETES','EPOC','ASMA','INMUSUPR','HIPERTENSION','OTRA_COM','CARDIOVASCULAR','OBESIDAD','RENAL_CRONICA']
dict_ill_pos = dict()
for name in ill_name:
    dict_ill_pos[name] = df_positive_cases.query(f'{name} == 1').shape[0]
print(f' Most common illness on total positive cases  = {max(dict_ill_pos, key=dict_ill_pos.get)} '.center(100,'#'))

# Data Visualization
graph_pie(dict_ill_pos,'Common illness on total positive cases')

##### Common illness on total positive cases by state

for state in sorted(df_cov.ENTIDAD_RES.unique()):
    print(f' {state} '.center(100,'#'))
    # Data Extraction
    ill_name = ['DIABETES','EPOC','ASMA','INMUSUPR','HIPERTENSION','OTRA_COM','CARDIOVASCULAR','OBESIDAD','RENAL_CRONICA']
    dict_ill_pos = dict()
    for name in ill_name:
        dict_ill_pos[name] = df_positive_cases.query(f'{name} == 1 & ENTIDAD_RES == "{state}"').shape[0]
    print(f'Most common illness on total positive cases at {state} = {max(dict_ill_pos, key=dict_ill_pos.get)}')
    # Data Visualization
    graph_pie(dict_ill_pos,f'Common illness on total positive cases at {state}')

# ## Deceased COVID-19 cases

##### Analysis of total number of deceased cases
# FECHA_DEF has 9999-99-99 when a person COVID-19 confirmed. didn't die.

df_deceased_cases =  df_cov[df_cov['FECHA_DEF'] !='9999-99-99']

fig, axes = plt.subplots(1, 3, figsize=(25, 5))

# 1st graph - Gender of positive cases
axes[0].set_title('Gender of deceased cases')
plot = sns.barplot(x = ['Women','Men'],y = df_deceased_cases['SEXO'].value_counts().values, palette = 'pastel', ax = axes[0])
plot.set(ylim = (1200000,1400000), ylabel = None, yticklabels = [])
plot.tick_params(left=False)
autolabel(plot)

# 2nd graph - Age distribution of deceased cases
axes[1].set_title('Age distribution of deceased cases')
plot = sns.distplot(df_deceased_cases.EDAD, ax = axes[1])
plot.set(ylabel = None, yticklabels = [])
plot.tick_params(left=False)

# 3rd graph - Time between FECHA_SINTOMA and FECHA_DEF when a person die of COVID-19
axes[2].set_title('Days between been positive \n in COVID-19 and die')
data = df_cov['DIFERENCIA'].value_counts()
plot = sns.barplot(x = data.index[1:15], y = data.values[1:15],
                   palette = 'pastel', order = data.index[1:15], ax = axes[2])
plot.set(ylim = (None,10000), ylabel = None, yticklabels = [])
plot.tick_params(left=False)
autolabel(plot)

plt.tight_layout()
plt.show()

##### Total number of deceased cases by state

# Data Extraction
print(f' Total number of deceased cases = {df_deceased_cases.shape[0]} | % of total = {(df_deceased_cases.shape[0]/df_cov.shape[0]):.3}% '.center(100,'#'))
data = df_deceased_cases.groupby('ENTIDAD_RES').count()['CLASIFICACION_FINAL'].sort_values(ascending = False)

# Data Visualization
graph_bar(data, 'Total number of deceased cases by state')

##### Mean age of deceased cases by state

# Data Extraction
print(f' Total mean age of deceased cases = {(df_deceased_cases.EDAD.mean()):.3} '.center(100,'#'))
data = df_deceased_cases.groupby('ENTIDAD_RES').mean()['EDAD'].sort_values(ascending = False)

# Data Visualization
graph_bar(data, 'Mean age of deceased cases by state')

##### Common illness on total deceased cases

# Data Extraction
ill_name = ['DIABETES','EPOC','ASMA','INMUSUPR','HIPERTENSION','OTRA_COM','CARDIOVASCULAR','OBESIDAD','RENAL_CRONICA']
dict_ill_dec = dict()
for name in ill_name:
    dict_ill_dec[name] = df_deceased_cases.query(f'{name} == 1').shape[0]
print(f' Most common illness on total deceased cases  = {max(dict_ill_dec, key=dict_ill_dec.get)} '.center(100,'#'))

# Data Visualization
graph_pie(dict_ill_dec,'Common illness on total deceased cases')

##### Common illness on total positive cases by state

for state in sorted(df_cov.ENTIDAD_RES.unique()):
    print(f' {state} '.center(100,'#'))
    # Data Extraction
    ill_name = ['DIABETES','EPOC','ASMA','INMUSUPR','HIPERTENSION','OTRA_COM','CARDIOVASCULAR','OBESIDAD','RENAL_CRONICA']
    dict_ill_dec = dict()
    for name in ill_name:
        dict_ill_dec[name] = df_deceased_cases.query(f'{name} == 1 & ENTIDAD_RES == "{state}"').shape[0]
    print(f'Most common illness on total positive cases at {state} = {max(dict_ill_dec, key=dict_ill_dec.get)}')
    # Data Visualization
    graph_pie(dict_ill_dec,f'Common illness on total positive cases at {state}')

# # Interactive map

##### Replace the name of the states in the DataFrame with the json -- to avoid missing information

with open("mexico22.json") as f:
    data = json.load(f)

states_json = list()
for i in range(32):
    states_json.append(data['features'][i]['properties']['name'])

states_json = sorted(states_json)
states_df = sorted(df_cov.ENTIDAD_RES.unique())
print('In json the differences appear as:',sorted(set(states_json) - set(states_df)))
print('While in the DataFrame the differences appear as:',sorted(set(states_df) - set(states_json)))

dict_states = dict()
for json,df in zip(sorted(set(states_json) - set(states_df)),sorted(set(states_df) - set(states_json))):
    dict_states[df] = json

df_cov['ENTIDAD_RES'] = df_cov['ENTIDAD_RES'].replace(dict_states)

##### Number of confirmed cases and deceased cases on Mexico (dataset)

data = gpd.read_file("mexico22.json").sort_values('name',ascending = True).reset_index(drop = True)
data.rename(columns = {'name': "States"}, inplace=True)
data['Positives'] = df_cov[(df_cov['CLASIFICACION_FINAL'].isin([1,2,3])) & (~(df_cov['ENTIDAD_NAC'].isin([97,98,99])))].groupby('ENTIDAD_RES')                     .count()['CLASIFICACION_FINAL'].values
data['Deaths'] = df_cov[(~(df_cov['FECHA_DEF'] =='2001-01-01')) & (~(df_cov['ENTIDAD_NAC'].isin([97,98,99])))].groupby('ENTIDAD_NAC')                  .count()['CLASIFICACION_FINAL'].values
data

##### Number of confirmed cases on Mexico (Map)
# Check the json's path before running it

# Creation of the map
mexico_map = folium.Map(location=[23.634501, -102.552784], zoom_start=5.5, tiles = None)
folium.TileLayer('http://tile.stamen.com/watercolor/{z}/{x}/{y}.png', name = "Watercolor map", control = False, attr = "toner-bcg").add_to(mexico_map)
mexico_geo = r"mexico22.json"

# Adding confirmed cases layer
choropleth = folium.Choropleth(
    name='Confirmed cases',
    geo_data = mexico_geo,
    data = data,
    columns = ['States','Positives'],
    key_on = 'feature.properties.name',
    fill_color = 'YlOrRd',
    fill_opacity = 0.65,
    line_opacity = 0.5,
    threshold_scale = list(data['Positives'].quantile([0, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 1])),
    overlay= False)

for key in choropleth._children:
    if key.startswith('color_map'):
        del(choropleth._children[key])
choropleth.add_to(mexico_map)

# Adding deceased cases layer
choropleth = folium.Choropleth(
    name='Deceased cases',
    geo_data = mexico_geo,
    data = data,
    columns = ['States','Deaths'],
    key_on = 'feature.properties.name',
    fill_color = 'YlOrBr',
    fill_opacity = 0.65,
    line_opacity = 0.5,
    threshold_scale = list(data['Deaths'].quantile([0, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 1])),
    overlay= False)

for key in choropleth._children:
    if key.startswith('color_map'):
        del(choropleth._children[key])
choropleth.add_to(mexico_map)

# Adding pop-up tooltips
style_function = lambda x: {'fillColor': '#ffffff',
                            'color':'#000000',
                            'fillOpacity': 0.1,
                            'weight': 0.1}

highlight_function = lambda x: {'fillColor': '#000000',
                                'color':'#000000',
                                'fillOpacity': 0.50,
                                'weight': 0.1}

data_Geo = gpd.GeoDataFrame(data , geometry = data.geometry)

pop_up = folium.features.GeoJson(
    data_Geo,
    style_function = style_function,
    control = False,
    highlight_function = highlight_function,
    tooltip = folium.features.GeoJsonTooltip(
        fields=['States','Positives', 'Deaths'],
        aliases=['State: ','Number of COVID-19 confirmed cases: ', 'Number of COVID-19 deceased cases: '],
        style=("background-color: white; color: #333333; font-family: arial; font-size: 12px;")))
mexico_map.add_child(pop_up)
mexico_map.keep_in_front(pop_up)

# To control the layers
folium.LayerControl(collapsed=False).add_to(mexico_map)

mexico_map


"""
Data Organization II - Sorting a DataFrame
"""

import numpy as np
import pandas as pd
wine_reviews = pd.read_csv('../winemag-data-130k.csv')

# Using this DataFrame, sort the rows reverse alphabetically based on the 'country' column.

wine_ratings = wine_reviews[['title', 'country', 'rating', 'price']]

# Now sort the DataFrame's rows based on price from highest to lowest.


import numpy as np
import pandas as pd

from pandas import DataFrame as df         # DataFrame으로 작성해서 난 에러는 df로 바꾸기

# # 파일 읽어오기
# # csv_test = pd.read_csv("test_csv_file.csv")
# # print(csv_test)
# # #    ID LAST_NAME  AGE
# # # 0   1       KIM   30
# # # 1   2      CHOI   25
# # # 2   3       LEE   41
# # # 3   4      PARK   19
# # # 4   5       LIM   36
# # print(csv_test.shape)
# # # (5, 3)
# #
# # text_test = pd.read_csv("test_text_file.txt",sep="|")
# # print(text_test)
# # #    ID  A  B  C  D
# # # 0  C1  1  2  3  4
# # # 1  C2  5  6  7  8
# # # 2  C3  1  3  5  7
# # text_test = pd.read_csv("test_text_file.txt",sep="|",index_col='ID')               #index_col= : index로 사용하고자 하는 column name
# # print(text_test)
# # #     A  B  C  D
# # # ID
# # # C1  1  2  3  4
# # # C2  5  6  7  8
# # # C3  1  3  5  7
# # text_test = pd.read_csv("test_text_file.txt",sep="|",index_col=0)               #index_col= : index로 사용하고자 하는 column number
# # print(text_test)
# # #     A  B  C  D
# # # ID
# # # C1  1  2  3  4
# # # C2  5  6  7  8
# # # C3  1  3  5  7
# #
# # print("="*250)
# #
# # text_test = pd.read_csv("text_without_column_name.txt",sep="|")             # 첫줄을 헤더로 인식해 정보를 잃음
# # print(text_test)
# # #    C1  1  2  3  4
# # # 0  C2  5  6  7  8
# # # 1  C3  1  3  5  7
# # text_test = pd.read_csv("text_without_column_name.txt",sep="|",header=None)             # header=None : 헤더 없음을 인식
# # print(text_test)
# # #     0  1  2  3  4
# # # 0  C1  1  2  3  4
# # # 1  C2  5  6  7  8
# # # 2  C3  1  3  5  7
# # text_test = pd.read_csv("text_without_column_name.txt",sep="|",header=None,names=['ID','A','B','C','D'])             # names= : 헤더에 이름 부여
# # print(text_test)
# # #    ID  A  B  C  D
# # # 0  C1  1  2  3  4
# # # 1  C2  5  6  7  8
# # # 2  C3  1  3  5  7
# # text_test = pd.read_csv("text_without_column_name.txt",sep="|",header=None,names=['ID','A','B','C','D'],index_col='ID')
# # print(text_test)
# # #     A  B  C  D
# # # ID
# # # C1  1  2  3  4
# # # C2  5  6  7  8
# # # C3  1  3  5  7
#
# #분석 결과를 파일로 저장
#
# data = {
#     'id' : ['a1','a2','a3','a4','a5'],
#     'x1' : [1,2,3,4,5],
#     'x2' : [3.0,4.5,3.2,4.0,3.5]
# }
# #아래와 같이 작성하면 행번호가 함께 생성
# data_df = DataFrame(data)
# print(data_df)
# #    id  x1   x2
# # 0  a1   1  3.0
# # 1  a2   2  4.5
# # 2  a3   3  3.2
# # 3  a4   4  4.0
# # 4  a5   5  3.5
#
# data_df = DataFrame(data, index=['a1','a2','a3','a4','a5'])
# print(data_df)
#
# print("="*50)
#
# data_df_2 = data_df.reindex(['a1','a2','a3','a4','a5','a6'])
# print(data_df_2)
# #      id   x1   x2
# # a1   a1  1.0  3.0
# # a2   a2  2.0  4.5
# # a3   a3  3.0  3.2
# # a4   a4  4.0  4.0
# # a5   a5  5.0  3.5
# # a6  NaN  NaN  NaN
#
# # data_df_2.to_csv() : DataFrame을 csv로 저장할 때 사용
# data_df_2.to_csv('data_df_2.csv',sep=',',na_rep='NaN')

#DataFrame속성작업
#index : 행이름을 부여하는 속성 dataframe을 만들때 직접 지정 가능
#copy의 초기값 : false
df_1=df(data =np.arange(12).reshape(3,4))
print(df_1)
#    0  1   2   3
# 0  0  1   2   3
# 1  4  5   6   7
# 2  8  9  10  11
df_1=df(data =np.arange(12).reshape(3,4),index=['r0','r1','r2'])
print(df_1)
#     0  1   2   3
# r0  0  1   2   3
# r1  4  5   6   7
# r2  8  9  10  11
df_1=df(data =np.arange(12).reshape(3,4),index=['r0','r1','r2'],dtype=int,columns=['c0','c1','c2','c3'])
print(df_1)
#     c0  c1  c2  c3
# r0   0   1   2   3
# r1   4   5   6   7
# r2   8   9  10  11

#행과 열 전환
df_2=df(df_1.T)
print(df_2)
#     r0  r1  r2
# c0   0   4   8
# c1   1   5   9
# c2   2   6  10
# c3   3   7  11
print(df_1.T)
#     r0  r1  r2
# c0   0   4   8
# c1   1   5   9
# c2   2   6  10
# c3   3   7  11

print(df_1.axes)                        #axis : 축정보
# [Index(['r0', 'r1', 'r2'], dtype='object'), Index(['c0', 'c1', 'c2', 'c3'], dtype='object')]
print(df_1.dtypes)                      #data_type정보
# c0    int32
# c1    int32
# c2    int32
# c3    int32
# dtype: object
print(df_1.size)                        #데이터의 사이즈 확인
# 12

print(type(df_1))                       #자료구조
# <class 'pandas.core.frame.DataFrame'
print(df_1.values)                      #추출된 자료
# [[ 0  1  2  3]
#  [ 4  5  6  7]
#  [ 8  9 10 11]]
print(type(df_1.values))                #추출된 자료구조
# <class 'numpy.ndarray'>

df_2=df(data={
    'class_1' : ['a','a','b','b','c'],
    'var_1' : np.arange(5),
    'var_2' : np.random.rand(5)},
    index=['r0','r1','r2','r3','r4'])
print(df_2)
#    class_1  var_1     var_2
# r0       a      0  0.531274
# r1       a      1  0.175335
# r2       b      2  0.959343
# r3       b      3  0.442552
# r4       c      4  0.862076

df_2=df({                                                           # data = : 속성을 지정하지 않아도 기본값으로 지정됨
    'class_1' : ['a','a','b','b','c'],
    'var_1' : np.arange(5),
    'var_2' : np.random.rand(5)},
    index=['r0','r1','r2','r3','r4'])
print(df_2)
#    class_1  var_1     var_2
# r0       a      0  0.531274
# r1       a      1  0.175335
# r2       b      2  0.959343
# r3       b      3  0.442552
# r4       c      4  0.862076

#index 확인
print(df_2.index)
# Index(['r0', 'r1', 'r2', 'r3', 'r4'], dtype='object')
print(df_2.ix[2:])                      #index를 활용한 슬라이싱 : [2:] =>2번부터 마지막까지(행범위 참조)
#    class_1  var_1     var_2
# r2       b      2  0.614414
# r3       b      3  0.611415
# r4       c      4  0.213260
print(df_2.ix[2])
print(df_2.head(3))                                 # 앞에서 3번째까지 슬라이싱
#    class_1  var_1     var_2
# r0       a      0  0.346856
# r1       a      1  0.461321
# r2       b      2  0.079013
# print(df_2.tail(3))                                 # 뒤에서 3번째까지 슬라이싱
# #    class_1  var_1     var_2
# # r2       b      2  0.079013
# # r3       b      3  0.221049
# # r4       c      4  0.100560
#
# print(df_2.columns)
# # Index(['class_1', 'var_1', 'var_2'], dtype='object')
# print(df_2['class_1'])
# # r0    a
# # r1    a
# # r2    b
# # r3    b
# # r4    c
# # Name: class_1, dtype: object
#
# #class_1과 var_2 컬럼 출력
# print(df_2)
# print(df_2[['class_1','var_2']])            #r언어 : c()
# #     class_1     var_2
# # r0       a  0.333455
# # r1       a  0.527994
# # r2       b  0.003908
# # r3       b  0.513808
# # r4       c  0.119298
#
# idx = ['r0','r1','r2','r3','r4']
#
# df_1 = df({'c1' : np.arange(5),
#            'c2' : np.random.randn(5)},
#           index=idx)
# print(df_1)
# #     c1        c2
# # r0   0  0.421806
# # r1   1  1.287233
# # r2   2  0.028156
# # r3   3 -0.591526
# # r4   4  0.287460
#
# new_idx = ['r0','r1','r2','r5','r6']
# # df_1 = df_1.reindex(new_idx)
# # print(df_1)
# # #      c1        c2
# # # r0  0.0  0.475256
# # # r1  1.0  0.491309
# # # r2  2.0  0.864039
# # # r5  NaN       NaN
# # # r6  NaN       NaN
#
# # df_1 = df_1.reindex(new_idx,fill_value=0)
# # print(df_1)
# # #     c1        c2
# # # r0   0 -0.062061
# # # r1   1  0.132808
# # # r2   2 -0.049169
# # # r5   0  0.000000
# # # r6   0  0.000000
#
# # df_1 = df_1.reindex(new_idx,fill_value='missing')
# # print(df_1)
# # #          c1       c2
# # # r0        0  1.29382
# # # r1        1  -0.3048
# # # r2        2  1.52637
# # # r5  missing  missing
# # # r6  missing  missing
#
# df_1 = df_1.reindex(new_idx,fill_value='NA')                # fill_value= : 위에 내용 주석이여야 함(데이터가 없어야 해서)
# print(df_1)
# #     c1         c2
# # r0   0 -0.0427524
# # r1   1   -2.05168
# # r2   2    1.90008
# # r5  NA         NA
# # r6  NA         NA
#
#
# #시계열 data : 시간의 흐름에 따라 달라짐(내용 안에 시간의 변수가 있다.)
#
# #help(pd.date_range)
# # date_range(start=None, end=None, periods=None, freq=None, tz=None, normalize=False, name=None, closed=None, **kwargs)
# #     Return a fixed frequency DatetimeIndex.
# #
# #     Parameters
# #     ----------
# #     start : str or datetime-like, optional
# #         Left bound for generating dates.
# #     end : str or datetime-like, optional
# #         Right bound for generating dates.
# #     periods : integer, optional
# #         Number of periods to generate.
# #     freq : str or DateOffset, default 'D' (calendar daily)
# #         Frequency strings can have multiples, e.g. '5H'. See
# #         :ref:`here <timeseries.offset_aliases>` for a list of
# #         frequency aliases.
# #     tz : str or tzinfo, optional
# #         Time zone name for returning localized DatetimeIndex, for example
# #         'Asia/Hong_Kong'. By default, the resulting DatetimeIndex is
# #         timezone-naive.
# #     normalize : bool, default False
# #         Normalize start/end dates to midnight before generating date range.
# #     name : str, default None
# #         Name of the resulting DatetimeIndex.
# #     closed : {None, 'left', 'right'}, optional
# #         Make the interval closed with respect to the given frequency to
# #         the 'left', 'right', or both sides (None, the default).
# #     **kwargs
# #         For compatibility. Has no effect on the result.
# #
# #     Returns
# #     -------
# #     rng : DatetimeIndex
# #
# #     See Also
# #     --------
# #     pandas.DatetimeIndex : An immutable container for datetimes.
# #     pandas.timedelta_range : Return a fixed frequency TimedeltaIndex.
# #     pandas.period_range : Return a fixed frequency PeriodIndex.
# #     pandas.interval_range : Return a fixed frequency IntervalIndex.
# #
# #     Notes
# #     -----
# #     Of the four parameters ``start``, ``end``, ``periods``, and ``freq``,
# #     exactly three must be specified. If ``freq`` is omitted, the resulting
# #     ``DatetimeIndex`` will have ``periods`` linearly spaced elements between
# #     ``start`` and ``end`` (closed on both sides).
# #
# #     To learn more about the frequency strings, please see `this link
# #     <http://pandas.pydata.org/pandas-docs/stable/timeseries.html#offset-aliases>`__.
# #
# #     Examples
# #     --------
# #     **Specifying the values**
# #
# #     The next four examples generate the same `DatetimeIndex`, but vary
# #     the combination of `start`, `end` and `periods`.
# #
# #     Specify `start` and `end`, with the default daily frequency.
# #
# #     >>> pd.date_range(start='1/1/2018', end='1/08/2018')
# #     DatetimeIndex(['2018-01-01', '2018-01-02', '2018-01-03', '2018-01-04',
# #                    '2018-01-05', '2018-01-06', '2018-01-07', '2018-01-08'],
# #                   dtype='datetime64[ns]', freq='D')
# #
# #     Specify `start` and `periods`, the number of periods (days).
# #
# #     >>> pd.date_range(start='1/1/2018', periods=8)
# #     DatetimeIndex(['2018-01-01', '2018-01-02', '2018-01-03', '2018-01-04',
# #                    '2018-01-05', '2018-01-06', '2018-01-07', '2018-01-08'],
# #                   dtype='datetime64[ns]', freq='D')
# #
# #     Specify `end` and `periods`, the number of periods (days).
# #
# #     >>> pd.date_range(end='1/1/2018', periods=8)
# #     DatetimeIndex(['2017-12-25', '2017-12-26', '2017-12-27', '2017-12-28',
# #                    '2017-12-29', '2017-12-30', '2017-12-31', '2018-01-01'],
# #                   dtype='datetime64[ns]', freq='D')
# #
# #     Specify `start`, `end`, and `periods`; the frequency is generated
# #     automatically (linearly spaced).
# #
# #     >>> pd.date_range(start='2018-04-24', end='2018-04-27', periods=3)
# #     DatetimeIndex(['2018-04-24 00:00:00', '2018-04-25 12:00:00',
# #                    '2018-04-27 00:00:00'], freq=None)
# #
# #     **Other Parameters**
# #
# #     Changed the `freq` (frequency) to ``'M'`` (month end frequency).
# #
# #     >>> pd.date_range(start='1/1/2018', periods=5, freq='M')
# #     DatetimeIndex(['2018-01-31', '2018-02-28', '2018-03-31', '2018-04-30',
# #                    '2018-05-31'],
# #                   dtype='datetime64[ns]', freq='M')
# #
# #     Multiples are allowed
# #
# #     >>> pd.date_range(start='1/1/2018', periods=5, freq='3M')
# #     DatetimeIndex(['2018-01-31', '2018-04-30', '2018-07-31', '2018-10-31',
# #                    '2019-01-31'],
# #                   dtype='datetime64[ns]', freq='3M')
# #
# #     `freq` can also be specified as an Offset object.
# #
# #     >>> pd.date_range(start='1/1/2018', periods=5, freq=pd.offsets.MonthEnd(3))
# #     DatetimeIndex(['2018-01-31', '2018-04-30', '2018-07-31', '2018-10-31',
# #                    '2019-01-31'],
# #                   dtype='datetime64[ns]', freq='3M')
# #
# #     Specify `tz` to set the timezone.
# #
# #     >>> pd.date_range(start='1/1/2018', periods=5, tz='Asia/Tokyo')
# #     DatetimeIndex(['2018-01-01 00:00:00+09:00', '2018-01-02 00:00:00+09:00',
# #                    '2018-01-03 00:00:00+09:00', '2018-01-04 00:00:00+09:00',
# #                    '2018-01-05 00:00:00+09:00'],
# #                   dtype='datetime64[ns, Asia/Tokyo]', freq='D')
# #
# #     `closed` controls whether to include `start` and `end` that are on the
# #     boundary. The default includes boundary points on either end.
# #
# #     >>> pd.date_range(start='2017-01-01', end='2017-01-04', closed=None)
# #     DatetimeIndex(['2017-01-01', '2017-01-02', '2017-01-03', '2017-01-04'],
# #                   dtype='datetime64[ns]', freq='D')
# #
# #     Use ``closed='left'`` to exclude `end` if it falls on the boundary.
# #
# #     >>> pd.date_range(start='2017-01-01', end='2017-01-04', closed='left')
# #     DatetimeIndex(['2017-01-01', '2017-01-02', '2017-01-03'],
# #                   dtype='datetime64[ns]', freq='D')
# #
# #     Use ``closed='right'`` to exclude `start` if it falls on the boundary.
# #
# #     >>> pd.date_range(start='2017-01-01', end='2017-01-04', closed='right')
# #     DatetimeIndex(['2017-01-02', '2017-01-03', '2017-01-04'],
# #                   dtype='datetime64[ns]', freq='D')
#
#
# # print(pd.date_range("2018-9-10","2019-9-30",freq='MS'))          #http://pandas.pydata.org/pandas-docs/stable/genindex.html#D 에서 찾아서 확인
# # # DatetimeIndex(['2018-10-01', '2018-11-01', '2018-12-01', '2019-01-01',
# # #                '2019-02-01', '2019-03-01', '2019-04-01', '2019-05-01',
# # #                '2019-06-01', '2019-07-01', '2019-08-01', '2019-09-01'],
# # #               dtype='datetime64[ns]', freq='MS')
#
#
# date_idx = pd.date_range('09/10/2018',periods=10,freq='D')
# print(date_idx)
# # DatetimeIndex(['2018-09-10', '2018-09-11', '2018-09-12', '2018-09-13',
# #                '2018-09-14', '2018-09-15', '2018-09-16', '2018-09-17',
# #                '2018-09-18', '2018-09-19'],
# #               dtype='datetime64[ns]', freq='D')
# df_2 = df({"c1": [10,20,30,40,50,10,20,30,40,50]},
#     index=date_idx)
# print(df_2)
# #             c1
# # 2018-09-10  10
# # 2018-09-11  20
# # 2018-09-12  30
# # 2018-09-13  40
# # 2018-09-14  50
# # 2018-09-15  10
# # 2018-09-16  20
# # 2018-09-17  30
# # 2018-09-18  40
# # 2018-09-19  50
#
# date_idx_2 = pd.date_range('09/5/2018',periods=20,freq='D')         #date_idx보다 날짜를 더 늘림
# print(date_idx_2)
# # DatetimeIndex(['2018-09-05', '2018-09-06', '2018-09-07', '2018-09-08',
# #                '2018-09-09', '2018-09-10', '2018-09-11', '2018-09-12',
# #                '2018-09-13', '2018-09-14', '2018-09-15', '2018-09-16',
# #                '2018-09-17', '2018-09-18', '2018-09-19', '2018-09-20',
# #                '2018-09-21', '2018-09-22', '2018-09-23', '2018-09-24'],
# #               dtype='datetime64[ns]', freq='D')
# # df_2 = df_2.reindex(date_idx_2)         #날짜 추가로 설정이 없는 날짜는 NaN
# # print(df_2)
# # #               c1
# # # 2018-09-05   NaN
# # # 2018-09-06   NaN
# # # 2018-09-07   NaN
# # # 2018-09-08   NaN
# # # 2018-09-09   NaN
# # # 2018-09-10  10.0
# # # 2018-09-11  20.0
# # # 2018-09-12  30.0
# # # 2018-09-13  40.0
# # # 2018-09-14  50.0
# # # 2018-09-15  10.0
# # # 2018-09-16  20.0
# # # 2018-09-17  30.0
# # # 2018-09-18  40.0
# # # 2018-09-19  50.0
# # # 2018-09-20   NaN
# # # 2018-09-21   NaN
# # # 2018-09-22   NaN
# # # 2018-09-23   NaN
# # # 2018-09-24   NaN
#
# # df_2 = df_2.reindex(date_idx_2,method='ffill')                  #ffill : 바로 앞의 데이터를 복사
# # print(df_2)
# # #               c1
# # # 2018-09-05   NaN
# # # 2018-09-06   NaN
# # # 2018-09-07   NaN
# # # 2018-09-08   NaN
# # # 2018-09-09   NaN
# # # 2018-09-10  10.0
# # # 2018-09-11  20.0
# # # 2018-09-12  30.0
# # # 2018-09-13  40.0
# # # 2018-09-14  50.0
# # # 2018-09-15  10.0
# # # 2018-09-16  20.0
# # # 2018-09-17  30.0
# # # 2018-09-18  40.0
# # # 2018-09-19  50.0
# # # 2018-09-20  50.0
# # # 2018-09-21  50.0
# # # 2018-09-22  50.0
# # # 2018-09-23  50.0
# # # 2018-09-24  50.0
#
#
# df_2 = df_2.reindex(date_idx_2,method='bfill')                  #bfill : 바로 뒤의 데이터를 복사
# print(df_2)
# #               c1
# # 2018-09-05  10.0
# # 2018-09-06  10.0
# # 2018-09-07  10.0
# # 2018-09-08  10.0
# # 2018-09-09  10.0
# # 2018-09-10  10.0
# # 2018-09-11  20.0
# # 2018-09-12  30.0
# # 2018-09-13  40.0
# # 2018-09-14  50.0
# # 2018-09-15  10.0
# # 2018-09-16  20.0
# # 2018-09-17  30.0
# # 2018-09-18  40.0
# # 2018-09-19  50.0
# # 2018-09-20   NaN
# # 2018-09-21   NaN
# # 2018-09-22   NaN
# # 2018-09-23   NaN
# # 2018-09-24   NaN
#
#
#
#
#
#
#
#
#
#
#
#
#
#
#
#


import torch
import copy
import sys
import numpy as np
from utils import one_hot_encode, capsnet_testing_loss
from torch.autograd import Variable
from torch.backends import cudnn
from quantization_methods import *
from quantized_models import *

def quantized_test(model, num_classes, data_loader, quantization_function, quantization_bits,
                   quantization_bits_routing):
    """ Function to test the accuracy of the quantized models

        Args:
            model: pytorch model
            num_classes: number ot classes of the dataset
            data_loader: data loader of the test dataset
            quantization_function: quantization function of the quantization method to use
            quantization_bits: list, quantization bits for the activations
            quantization_bits_routing: list, quantization bits for the dynamic routing
        Returns:
            accuracy_percentage: accuracy of the quantized model expressed in percentage """
    # Switch to evaluate mode
    model.eval()

    loss = 0
    correct = 0

    num_batches = len(data_loader)

    for data, target in data_loader:
        batch_size = data.size(0)
        target_one_hot = one_hot_encode(target, length=num_classes)

        if torch.cuda.device_count() > 0:  # if there are available GPUs, move data to the first visible
            device = torch.device("cuda:0")
            data = data.to(device)
            target = target.to(device)
            target_one_hot = target_one_hot.to(device)

        # Output predictions
        output = model(data, quantization_function, quantization_bits, quantization_bits_routing)

        # Sum up batch loss
        m_loss = \
            capsnet_testing_loss(output, target_one_hot)
        loss += m_loss.data

        # Count number of correct predictions
        # Compute the norm of the vector capsules
        v_length = torch.sqrt((output ** 2).sum(dim=2))
        assert v_length.size() == torch.Size([batch_size, num_classes])

        # Find the index of the longest vector
        _, max_index = v_length.max(dim=1)
        assert max_index.size() == torch.Size([batch_size])

        # vector with 1 where the model makes a correct prediction, 0 where false
        correct_pred = torch.eq(target.cpu(), max_index.data.cpu())
        correct += correct_pred.sum()

    # Log test accuracies
    num_test_data = len(data_loader.dataset)
    accuracy_percentage = float(correct) * 100.0 / float(num_test_data)

    return accuracy_percentage

def qcapsnets(model, model_parameters, full_precision_filename, num_classes, data_loader, top_accuracy,
              accuracy_tolerance, memory_budget, quantization_scheme):
    """ Q-CapsNets framework - Quantization

        Args:
            model: string, name of the model
            model_parameters: list, parameters to use for the instantiation of the model class
            full_precision_filename: string, directory of the full-precision weights
            num_classes: number of classes of the dataset
            data_loader: data loader of the testing dataset
            top_accuracy : maximum accuracy reached by the full_precision trained model (percentage)
            accuracy_tolerance: tolerance of the quantized model accuracy with respect to the full precision accuracy.
                                Provided in percentage
            memory_budget: memory budget for the weights of the model. Provided in MB (MegaBytes)
            quantization_scheme: quantization scheme to be used by the framework (string, e.g., "truncation)"
        Returns:
            void
    """
    print("==> Q-CapsNets Framework")
    # instantiate the quantized model with the full-precision weights
    model_quant_class = getattr(sys.modules[__name__], model)
    model_quant_original = model_quant_class(*model_parameters)
    model_quant_original.load_state_dict(torch.load(full_precision_filename))

    # Move the model to GPU if available
    if torch.cuda.device_count() > 0:
        device = torch.device("cuda:0")
        model_quant_original.to(device)
        cudnn.benchmark = True

    # create the quantization functions
    possible_functions = globals().copy()
    possible_functions.update(locals())
    quantization_function_activations = possible_functions.get(quantization_scheme)
    if not quantization_function_activations:
        raise NotImplementedError("Quantization function %s not implemented" % quantization_scheme)
    quantization_function_weights = possible_functions.get(quantization_scheme + "_inplace")
    if not quantization_function_weights:
        raise NotImplementedError("Quantization function %s not implemented (inplace version)" % quantization_scheme)

    # compute the accuracy reduction available for each step
    minimum_accuracy = top_accuracy - accuracy_tolerance / 100 * top_accuracy
    acc_reduction = top_accuracy - minimum_accuracy
    step1_reduction = 5 / 100 * acc_reduction
    step1_min_acc = top_accuracy - step1_reduction

    print("Full-precision accuracy: ", top_accuracy, "%")
    print("Minimum quantized accuracy: ", minimum_accuracy, "%")
    print("Memory budget: ", memory_budget, "MB")
    print("Quantization method: ", quantization_scheme)
    print("\n")

    # STEP 1: Layer-Uniform quantization of weights and activations
    print("STEP 1")

    def step1_quantization_test(quantization_bits):
        """ Function to test the model at STEP 1 of the algorithm

            The function receives a single "quantization_bits" value N, and creates two lists [N, N, ..., N] and
            [N, N, ..., N] for the activations and the dynamic routing, since at STEP 1 all the layers are quantized
            uniformly. The weights of each layer are quantized with N bits too and then the accuracy of the model
            is computed.

            Args:
                quantization_bits: single value used for quantizing all the weights and activations
            Returns:
                acc_temp: accuracy of the model quantized uniformly with quantization_bits bits
        """
        quantized_model_temp = copy.deepcopy(model_quant_original)

        step1_act_bits_f = []     # list with the quantization bits for the activations
        step1_dr_bits_f = []      # list with the quantization bits for the dynamic routing
        for c in quantized_model_temp.children():
            step1_act_bits_f.append(quantization_bits)
            if c.capsule_layer:
                if c.dynamic_routing:
                    step1_dr_bits_f.append(quantization_bits)
            for p in c.parameters():
                with torch.no_grad():
                    quantization_function_weights(p, quantization_bits)      # Quantize the weights
        # test with quantized weights and activations
        acc_temp = quantized_test(quantized_model_temp, num_classes, data_loader,
                                  quantization_function_activations, step1_act_bits_f, step1_dr_bits_f)
        del quantized_model_temp
        return acc_temp

    # BINARY SEARCH of the bitwidth for step 1, starting from 32 bits
    step1_bit_search = [32]
    step1_acc_list = []      # list of accuracy at each step of the search algorithm
    step1_acc = step1_quantization_test(32)
    step1_acc_list.append(step1_acc)
    if step1_acc > step1_min_acc:
        step1_bit_search_sat = [True]    # True is the accuracy is higher than the minimum required
        step1_bit_search.append(16)
        while True:
            step1_acc = step1_quantization_test(step1_bit_search[-1])
            step1_acc_list.append(step1_acc)
            if step1_acc > step1_min_acc:
                step1_bit_search_sat.append(True)
            else:
                step1_bit_search_sat.append(False)
            if (abs(step1_bit_search[-1] - step1_bit_search[-2])) == 1:
                step1_bit_search_sat.reverse()
                step1_bits = step1_bit_search[
                    len(step1_bit_search_sat) - 1 - next(k for k, val in enumerate(step1_bit_search_sat) if val)]
                step1_bit_search_sat.reverse()
                step1_acc = step1_acc_list[
                    len(step1_bit_search_sat) - 1 - next(k for k, val in enumerate(step1_bit_search_sat) if val)]
                break
            else:
                if step1_acc > step1_min_acc:
                    step1_bit_search.append(
                        int(step1_bit_search[-1] - abs(step1_bit_search[-1] - step1_bit_search[-2]) / 2))
                else:
                    step1_bit_search.append(
                        int(step1_bit_search[-1] + abs(step1_bit_search[-1] - step1_bit_search[-2]) / 2))
    else:
        step1_bits = 32
        step1_acc = step1_acc_list[1]

    # Create the lists of bits ofSTEP 1
    step1_act_bits = []
    step1_dr_bits = []
    step1_weight_bits = []
    for c in model_quant_original.children():
        step1_act_bits.append(step1_bits)
        step1_weight_bits.append(step1_bits)
        if c.capsule_layer:
            if c.dynamic_routing:
                step1_dr_bits.append(step1_bits)

    print("STEP 1 output: ")
    print("\t Weight bits: \t\t", step1_weight_bits)
    print("\t Activation bits: \t\t", step1_act_bits)
    print("\t Dynamic Routing bits: \t\t", step1_dr_bits)
    print("STEP 1 accuracy: ", step1_acc)
    print("\n")

    # STEP2 - satisfy memory requirement
    # compute the number of weights and biases of each layer/block
    print("STEP 2")
    number_of_weights_inlayers = []
    for c in model_quant_original.children():
        param_intra_layer = 0
        for p in c.parameters():
            param_intra_layer = param_intra_layer + p.numel()
        number_of_weights_inlayers.append(param_intra_layer)
    number_of_blocks = len(number_of_weights_inlayers)

    memory_budget_bits = memory_budget * 8000000      # From MB to bits
    minimum_mem_required = np.sum(number_of_weights_inlayers)

    if memory_budget_bits < minimum_mem_required:
        raise ValueError("The memory budget can not be satisfied, increase it to",
                         minimum_mem_required / 8000000, " MB at least")

    # Compute the number of bits that satisfy the memory budget.
    # First try with [N, N-1, N-2, N-3, N-4, N-4, ...].
    # If it is not possible, try with [N, N-1, N-2, N-3, N-3, ...]
    # and so on until [N, N, N, N, ...] (number of bits uniform across the layers)
    decrease_amount = 5
    while decrease_amount >= 0:
        bit_decrease = []
        if number_of_blocks <= decrease_amount:
            i = 0
            for r in range(0, number_of_blocks):
                bit_decrease.append(i)
                i = i - 1
        else:
            i = 0
            for r in range(0, decrease_amount):
                bit_decrease.append(i)
                i = i - 1
            for r in range(decrease_amount, number_of_blocks):
                bit_decrease.append(i + 1)

        bits_memory_sat = 33
        while True:
            # decrease N (bits_memory_sat) until the memory budget is satisfied.
            bits_memory_sat = bits_memory_sat - 1
            memory_occupied = np.sum(np.multiply(number_of_weights_inlayers, np.add(bits_memory_sat + 1, bit_decrease)))
            # +1 because bits_memory_sat are the fractional part bits, but we need one for the integer part
            if memory_occupied <= memory_budget_bits:
                break

        step2_weight_bits = list(np.add(bits_memory_sat, bit_decrease))
        if step2_weight_bits[-1] >= 0:
            break
        else:
            decrease_amount = decrease_amount - 1

    # lists of bitwidths for activations and dynamic routing at STEP 1
    step2_act_bits = copy.deepcopy(step1_act_bits)
    step2_dr_bits = copy.deepcopy(step1_dr_bits)

    # Quantizeed the weights
    model_memory = copy.deepcopy(model_quant_original)
    for i, c in enumerate(model_memory.children()):
        for p in c.parameters():
            with torch.no_grad():
                quantization_function_weights(p, step2_weight_bits[i])
    step2_acc = quantized_test(model_memory, num_classes, data_loader,
                               quantization_function_activations, step2_act_bits, step2_dr_bits)

    print("STEP 2 output: ")
    print("\t Weight bits: \t\t", step2_weight_bits)
    print("\t Activation bits: \t\t", step2_act_bits)
    print("\t Dynamic Routing bits: \t\t", step2_dr_bits)
    print("STEP 2 accuracy: ", step2_acc)
    print("\n")

    # IF the step 2 accuracy is higher that the minimum required accuracy --> BRANCH A
    if step2_acc > minimum_accuracy:
        # What is the accuracy that can still be consumed?
        branchA_accuracy_budget = step2_acc - minimum_accuracy
        step3A_min_acc = step2_acc - branchA_accuracy_budget * 55 / 100

        # STEP 3A  - layer-wise quantization of activations
        print("STEP 3A")
        # get the position of the layers that use dynamic routing bits
        dynamic_routing_bits_bool = []
        for c in model_memory.children():
            if c.capsule_layer:
                if c.dynamic_routing:
                    dynamic_routing_bits_bool.append(True)
            else:
                dynamic_routing_bits_bool.append(False)
        layers_dr_position = [pos for pos, val in enumerate(dynamic_routing_bits_bool) if val]

        step3a_weight_bits = copy.deepcopy(step2_weight_bits)
        step3a_act_bits = copy.deepcopy(step2_act_bits)
        step3a_dr_bits = copy.deepcopy(step2_dr_bits)
        for l in range(0, len(step3a_act_bits)):
            while True:
                step3a_acc = quantized_test(model_memory, num_classes, data_loader,
                                            quantization_function_activations, step3a_act_bits, step3a_dr_bits)
                if step3a_acc >= step3A_min_acc:
                    step3a_act_bits[l:] = list(np.add(step3a_act_bits[l:], -1))
                    for x in range(len(layers_dr_position)):
                        step3a_dr_bits[x] = step3a_act_bits[layers_dr_position[x]]
                else:
                    step3a_act_bits[l:] = list(np.add(step3a_act_bits[l:], +1))
                    for x in range(len(layers_dr_position)):
                        step3a_dr_bits[x] = step3a_act_bits[layers_dr_position[x]]
                    break

        step3a_acc = quantized_test(model_memory, num_classes, data_loader,
                                    quantization_function_activations, step3a_act_bits, step3a_dr_bits)

        print("STEP 3A output: ")
        print("\t Weight bits: \t\t", step3a_weight_bits)
        print("\t Activation bits: \t\t", step3a_act_bits)
        print("\t Dynamic Routing bits: \t\t", step3a_dr_bits)
        print("STEP 3A accuracy: ", step3a_acc)
        print("\n")

        # STEP 4A  -  layer-wise quantization of dynamic routing
        print("STEP 4A")
        step4a_weight_bits = copy.deepcopy(step2_weight_bits)
        step4a_act_bits = copy.deepcopy(step3a_act_bits)
        step4a_dr_bits = copy.deepcopy(step3a_dr_bits)

        # need to variate only the bits of the layers in which the dynamic routing is actually performed
        # (iterations > 1)
        dynamic_routing_quantization = []
        for c in model_memory.children():
            if c.capsule_layer:
                if c.dynamic_routing:
                    if c.dynamic_routing_quantization:
                        dynamic_routing_quantization.append(True)
                    else:
                        dynamic_routing_quantization.append(False)
        dr_quantization_pos = [pos for pos, val in enumerate(dynamic_routing_quantization) if val]

        # new set of bits only if dynamic routing is performed
        dr_quantization_bits = [step4a_dr_bits[x] for x in dr_quantization_pos]
        for l in range(0, len(dr_quantization_bits)):
            while True:
                step4a_acc = quantized_test(model_memory, num_classes, data_loader,
                                            quantization_function_activations, step4a_act_bits, step4a_dr_bits)
                if step4a_acc >= minimum_accuracy:
                    dr_quantization_bits[l:] = list(np.add(dr_quantization_bits[l:], -1))
                    # update the whole vector step4a_dr_bits
                    for x in range(0, len(dr_quantization_bits)):
                        step4a_dr_bits[dr_quantization_pos[x]] = dr_quantization_bits[x]
                else:
                    dr_quantization_bits[l:] = list(np.add(dr_quantization_bits[l:], +1))
                    # update the whole vector step4a_dr_bits
                    for x in range(0, len(dr_quantization_bits)):
                        step4a_dr_bits[dr_quantization_pos[x]] = dr_quantization_bits[x]
                    break

        step4a_acc = quantized_test(model_memory, num_classes, data_loader,
                                    quantization_function_activations, step4a_act_bits, step4a_dr_bits)

        print("STEP 4A output: ")
        print("\t Weight bits: \t\t", step4a_weight_bits)
        print("\t Activation bits: \t\t", step4a_act_bits)
        print("\t Dynamic Routing bits: \t\t", step4a_dr_bits)
        print("STEP 4A accuracy: ", step4a_acc)
        print("\n")

        print("\n")
        quantized_filename = full_precision_filename[:-3] + '_quantized_satisfied.pt'
        torch.save(model_memory.state_dict(), quantized_filename)
        print("Model-satisfied stored in ", quantized_filename)
        print("\t Weight bits: \t\t", step4a_weight_bits)
        print("\t Activation bits: \t\t", step4a_act_bits)
        print("\t Dynamic Routing bits: \t\t", step4a_dr_bits)
        print("Model-satisfied accuracy: ", step4a_acc)

    else:
        # BRANCH B - STEP 3B  - layer-wise quantization of the weights
        print("STEP 3B")
        step3b_weight_bits = copy.deepcopy(step1_weight_bits)
        step3b_act_bits = copy.deepcopy(step1_act_bits)
        step3b_dr_bits = copy.deepcopy(step1_dr_bits)

        model_accuracy = copy.deepcopy(model_quant_original)
        for i, c in enumerate(model_accuracy.children()):
            for p in c.parameters():
                with torch.no_grad():
                    quantization_function_weights(p, step3b_weight_bits[i])

        for l in range(0, len(step3b_weight_bits)):
            while True:
                step3b_acc = quantized_test(model_accuracy, num_classes, data_loader,
                                            quantization_function_activations, step3b_act_bits, step3b_dr_bits)
                if step3b_acc >= minimum_accuracy:
                    step3b_weight_bits[l:] = list(np.add(step3b_weight_bits[l:], -1))
                    model_accuracy = copy.deepcopy(model_quant_original)
                    for i, c in enumerate(model_accuracy.children()):
                        for p in c.parameters():
                            with torch.no_grad():
                                quantization_function_weights(p, step3b_weight_bits[i])
                else:
                    step3b_weight_bits[l:] = list(np.add(step3b_weight_bits[l:], +1))
                    model_accuracy = copy.deepcopy(model_quant_original)
                    for i, c in enumerate(model_accuracy.children()):
                        for p in c.parameters():
                            with torch.no_grad():
                                quantization_function_weights(p, step3b_weight_bits[i])
                    break

        step3b_acc = quantized_test(model_accuracy, num_classes, data_loader,
                                    quantization_function_activations, step3b_act_bits, step3b_dr_bits)

        print("STEP 3B output: ")
        print("\t Weight bits: \t\t", step3b_weight_bits)
        print("\t Activation bits: \t\t", step3b_act_bits)
        print("\t Dynamic Routing bits: \t\t", step3b_dr_bits)
        print("STEP 3B accuracy: ", step3b_acc)
        print("\n")

        print("\n")
        quantized_filename = full_precision_filename[:-3] + '_quantized_memory.pt'
        torch.save(model_memory.state_dict(), quantized_filename)
        print("Model-memory stored in ", quantized_filename)
        print("\t Weight bits: \t\t", step2_weight_bits)
        print("\t Activation bits: \t\t", step2_act_bits)
        print("\t Dynamic Routing bits: \t\t", step2_dr_bits)
        print("Model_memory accuracy: ", step2_acc)
        print("\n")
        quantized_filename = full_precision_filename[:-3] + '_quantized_accuracy.pt'
        torch.save(model_accuracy.state_dict(), quantized_filename)
        print("Model-memory stored in ", quantized_filename)
        print("\t Weight bits: \t\t", step3b_weight_bits)
        print("\t Activation bits: \t\t", step3b_act_bits)
        print("\t Dynamic Routing bits: \t\t", step3b_dr_bits)
        print("Model_accuracy accuracy: ", step3b_acc)


import cv2 as cv
import numpy as np

# 遍历读取图片的每一个像素点，打印出其属性值
# def access_pixels(image):
    # print(image.shape)
    # height = image.shape[0]
    # width = image.shape[1]
    # channels = image.shape[2]
    # print("width:%s,height:%s,channels:%s"% (width,height,channels))

    # 遍历每一个像素点，并且修改一个像素点的值。
    # for row in range(height):
    #     for col in range(width):
    #         for c in range(channels):
    #             pv = image[row,col,c]
    #             image[row,col,c] = 255 - pv
    # cv.imshow("P1",image)

# 当前所用的时间为：458.1146 ms

# 上面的操作就是对像素点进行取反操作。
# 另外就是Open CV定义了专门的API->bitwise_not()函数进行对像素点进行取反操作
# 因为Open CV是用C语言来写的。所以，在运算起来的时候就是比较快的
# 所以在运用的时候，我们就要更加高效率的使用Open CV内置的函数，可以提高效率
def inverse(image):
    dst = cv.bitwise_not(image)
    cv.imshow("Inverse Windows06",dst)
# # 运行速度：当前所用的时间为：32.9174 ms

# 创建一张图片，定义图片的大小以及dtype类型
def create_image():

    # 创建一个多通道的8位图像进行显示，
    # 修改图片的像素点的值
    # 创建一个三通道的长宽为400，并且dtype类型为np.uint8
    img = np.zeros([400,400,3],np.uint8)

#     但是我们同样可以对创建的图片的每一个像素点进行修改，
#     对第一个通道的像素点进行修改，
    img[:,:,0] = np.ones([400,400])*255

#     此时就已经可以进行显示了
    cv.imshow('Image Windows02',img)

    #  2.创建一个单通道8位的图像进行显示
    # 在这里我们学到了怎么初始化图像，
    # 一种是可以将图像得每一个的像素点都变成0
    # 另外一种就是将所有的像素点都变成1，并且可以进行乘法改变像素的大小
    # 单通道的一定也要加上通道数1,要不然会认为是二维的。
    # 并且会报错，IndexError: too many indices for array

    # 1. 初始化的时候，使用0,进行初始化
    # img = np.zeros([400,400,1],np.uint8)
    # # 127就是灰度图像
    # img[:, :, 0] = np.ones([400,400])*127

    # 2. 初始化的时候使用1进行初始化，会更加灵活一点
    img = np.ones([400, 400, 1], np.uint8)

    # 得到的是灰度图像
    img1 = img*127

    # 得到的是黑色图像
    img2 = img*0

    # 显示图像
    cv.imshow("Image Windows04",img1)
    cv.imshow("Image Windows05",img2)

    # 保存图像到本地dist/images/turtor_01/img_127.jpg和img_0.jpg
    cv.imwrite('./dist/images/turtor_01/img_127.jpg',img1)
    cv.imwrite('./dist/images/turtor_01/img_0.jpg',img2)

    # cv.imshow("Image windows03",img)

    # 3. 初始化一个多维数组,在以后进行图像处理时，如果有浮点型的，一定要用浮点型
    # 避免数据被截断
    m1 = np.ones([3,3],np.float)
    m1.fill(12.33)
    print(m1)

#     使用reshape()函数进行转换数组的维度
#     reshape()函数只是改变的数据的表示形式，但是并不会改变数据
    m2 = m1.reshape([1,9])
    print(m2)

print("----------- Hello Guyan -----------")
src = cv.imread('./images/handsomeboy01.jpg')
cv.namedWindow("Image Windows01",cv.WINDOW_AUTOSIZE)
cv.imshow("Image Windows01",src)

# getTickCount()返回从操作F系统启动到当前所经过的计时周期数
t1 = cv.getTickCount()
# access_pixels(src)
# getTickCount()函数，返回从操作系统启动到当前的计时周期数

# 此时就只是调用自己创建图片的函数create_image()
# create_image()
inverse(src)

t2 = cv.getTickCount()
# getTickFrequency()函数，用于返回CPU的频率，
# (getTickCount1 - getTickCount2)/getTickFrequency():
# (当前次数 - 开始计时次数) / 每秒钟重复次数 = 等于从开始到当前所用的时间。
time = (t2 - t1) / cv.getTickFrequency()

# 注意这里一定要带上括号，要不然就会将当前是所用的时间打印1000次
# print('当前所用的时间为：%s ms' % time*1000)
print('当前所用的时间为：%s ms' % (time*1000))
cv.waitKey(0)
cv.destroyAllWindows()


import os
import sys
import unittest
import numpy
from os.path import join as pjn

import QENSmodels

# resolve path to reference_data
this_module_path = sys.modules[__name__].__file__
data_dir = pjn(os.path.dirname(this_module_path), 'reference_data')

class TestBackgroundPolynomials(unittest.TestCase):
    """ Tests QENSmodels.background_polynomials function"""

    def test_type_output(self):
        """ test types of outputs depending on types of inputs of x
        (float or array)"""
        self.assertIsInstance(QENSmodels.background_polynomials(1, [1, 2, 3]),
                              numpy.float64)

        self.assertIsInstance(QENSmodels.background_polynomials([1, 2, 3],
                                                                [1, 2, 3]),
                              numpy.ndarray)

    def test_output_when_no_coeff(self):
        """ test that output = 0 if no list of coefficients given """
        self.assertEqual(QENSmodels.background_polynomials(1), 0.0)
        self.assertEqual(QENSmodels.background_polynomials(3.21), 0.0)

    def test_raised_error(self):
        """ test that an exception is raised if the input list of coefficients
        is incorrect, for example, if not all elements are numbers"""
        self.assertRaises(ValueError, QENSmodels.background_polynomials,
                          [1, 2, 3],
                          [1, 2, 'a'])

    def test_reference_data(self):
        """ Test output values in comparison with reference data
                   (file in 'reference data' folder) """

        # load reference data
        ref_data = numpy.loadtxt(pjn(data_dir,
                                     "background_polynomials_ref_data.dat"))

        # generate data from current model
        # for info: the parameters' values used for the reference data are
        # specified in the README file in the 'reference data' folder
        w = numpy.arange(-2, 2.01, 0.01)
        actual_data = numpy.column_stack([w, QENSmodels.background_polynomials(w, [1, 2, 3])]) # noqa:

        # compare the 2 arrays
        numpy.testing.assert_array_almost_equal(ref_data,
                                                actual_data,
                                                decimal=13)

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


# -*- encoding: utf-8 -*-
"""
@Author     :   zYx.Tom
@Contact    :   526614962@qq.com
@site       :   https://github.com/zhuyuanxiang/tensorflow_cookbook
---------------------------
@Software   :   PyCharm
@Project    :   TensorFlow_Machine_Learning_Cookbook
@File       :   tools.py
@Version    :   v0.1
@Time       :   2019-11-02 11:58
@License    :   (C)Copyright 2018-2019, zYx.Tom
@Reference  :   《TensorFlow机器学习实战指南，Nick McClure》, Sec04，P
@Desc       :   基于 TensorFlow 的线性回归，常用的 Python 工具函数
"""
# common imports
import os
import sys

import numpy as np  # pip install numpy<1.17，小于1.17就不会报错
import sklearn
import tensorflow as tf
from tensorflow.python.framework import ops

# 设置数据显示的精确度为小数点后3位
np.set_printoptions(precision = 8, suppress = True, threshold = np.inf, linewidth = 200)

# 利用随机种子，保证随机数据的稳定性，使得每次随机测试的结果一样
np.random.seed(42)

# 初始化默认的计算图
ops.reset_default_graph()
# Python ≥3.5 is required
assert sys.version_info >= (3, 5)
# Scikit-Learn ≥0.20 is required
assert sklearn.__version__ >= "0.20"
# 屏蔽警告：Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX2 FMA
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
# Open graph session
sess = tf.Session()

def show_values(variable, title = None, feed_dict = None, session = None):
    if type(title) is not str:
        print("Show_values()函数被重构了，把变量放在了第一个参数，把标题放在了第二个参数")
    if title is None:
        title = str(variable)
    if session is None:
        session = tf.Session()
        session.run(tf.global_variables_initializer())
    print('-' * 50)
    if title is not None:
        print("{} = {}".format(title, variable))
        print("session.run({}) = ".format(variable))
    result = session.run(variable, feed_dict = feed_dict)
    print(result)
    return result

def show_title(number_title):
    print('\n', '-' * 5, number_title, '-' * 5)