import dataclasses
import datetime
import io
import pathlib
import pickle

import pytest
import pandas as pd
import numpy as np
import structlog

from datapublic.common_fields import CommonFields
from datapublic.common_fields import DemographicBucket
from datapublic.common_fields import FieldName
from datapublic.common_fields import PdFields

from datapublic.common_test_helpers import to_dict

from libs import github_utils
from libs.datasets import AggregationLevel
from libs.datasets import dataset_pointer
from libs.datasets import taglib

from libs.datasets import timeseries
from libs.datasets.taglib import TagType
from libs.datasets.taglib import UrlStr
from libs.pipeline import Region
from tests import test_helpers
from tests.dataset_utils_test import read_csv_and_index_fips
from tests.dataset_utils_test import read_csv_and_index_fips_date
from tests.test_helpers import TimeseriesLiteral


# turns all warnings into errors for this module
pytestmark = pytest.mark.filterwarnings("error", "ignore::libs.pipeline.BadFipsWarning")


def _make_dataset_pointer(tmpdir, filename: str = "somefile.csv") -> dataset_pointer.DatasetPointer:
    # The fixture passes in a py.path, which is not the type in DatasetPointer.
    path = pathlib.Path(tmpdir) / filename

    fake_git_summary = github_utils.GitSummary(sha="abcdef", branch="main", is_dirty=True)

    return dataset_pointer.DatasetPointer(
        dataset_type=dataset_pointer.DatasetType.MULTI_REGION,
        path=path,
        model_git_info=fake_git_summary,
        updated_at=datetime.datetime.utcnow(),
    )


@pytest.mark.parametrize("include_na_at_end", [False, True])
def test_remove_padded_nans(include_na_at_end):
    rows = [
        {"date": "2020-02-01", "cases": pd.NA},
        {"date": "2020-02-02", "cases": pd.NA},
        {"date": "2020-02-03", "cases": 1},
        {"date": "2020-02-04", "cases": pd.NA},
        {"date": "2020-02-05", "cases": 2},
        {"date": "2020-02-06", "cases": 3},
    ]
    if include_na_at_end:
        rows += [{"date": "2020-02-07", "cases": pd.NA}]

    df = pd.DataFrame(rows)

    results = timeseries._remove_padded_nans(df, ["cases"])
    expected_series = pd.Series([1, pd.NA, 2, 3], name="cases")

    pd.testing.assert_series_equal(results.cases, expected_series)


def test_multi_region_to_from_timeseries_and_latest_values(tmp_path: pathlib.Path):
    # TODO(tom): Replace csv with test_helpers builders and uncomment assert in add_fips_static_df
    ts_df = read_csv_and_index_fips_date(
        "fips,county,aggregate_level,date,m1,m2\n"
        "97111,Bar County,county,2020-04-02,2,\n"
        "97222,Foo County,county,2020-04-01,,10\n"
        "01,,state,2020-04-01,,20\n"
    ).reset_index()
    latest_values_df = read_csv_and_index_fips(
        "fips,county,aggregate_level,c1,c2\n"
        "97111,Bar County,county,3,\n"
        "97222,Foo County,county,4,10.5\n"
        "01,,state,,123.4\n"
    ).reset_index()
    multiregion = (
        timeseries.MultiRegionDataset.from_fips_timeseries_df(ts_df)
        .add_fips_static_df(latest_values_df)
        .add_provenance_csv(
            io.StringIO("location_id,variable,provenance\n" "iso1:us#fips:97111,m1,ts197111prov\n")
        )
    )
    region_97111 = multiregion.get_one_region(Region.from_fips("97111"))
    assert region_97111.date_indexed.at["2020-04-02", "m1"] == 2
    assert region_97111.latest["c1"] == 3
    assert multiregion.get_one_region(Region.from_fips("01")).latest["c2"] == 123.4

    csv_path = tmp_path / "multiregion.csv"
    multiregion.to_csv(csv_path)
    multiregion_loaded = timeseries.MultiRegionDataset.from_csv(csv_path)
    region_97111 = multiregion_loaded.get_one_region(Region.from_fips("97111"))
    assert region_97111.date_indexed.at["2020-04-02", "m1"] == 2
    assert region_97111.latest["c1"] == 3
    assert region_97111.region.fips == "97111"
    assert multiregion_loaded.get_one_region(Region.from_fips("01")).latest["c2"] == 123.4
    test_helpers.assert_dataset_like(
        multiregion, multiregion_loaded, drop_na_latest=True, drop_na_timeseries=True
    )


def test_multi_region_get_one_region():
    ts = timeseries.MultiRegionDataset.from_csv(
        io.StringIO(
            "location_id,county,aggregate_level,date,m1,m2\n"
            "iso1:us#fips:97111,Bar County,county,2020-04-02,2,\n"
            "iso1:us#fips:97222,Foo County,county,2020-04-01,,10\n"
            "iso1:us#fips:97111,Bar County,county,,3,\n"
            "iso1:us#fips:97222,Foo County,county,,,11\n"
        )
    )
    region_97111_ts = ts.get_one_region(Region.from_fips("97111"))
    assert to_dict(["date"], region_97111_ts.data[["date", "m1", "m2"]]) == {
        pd.to_datetime("2020-04-02"): {"m1": 2}
    }
    assert region_97111_ts.latest["m1"] == 3
    assert region_97111_ts.region.fips == "97111"

    region_97222_ts = ts.get_one_region(Region.from_fips("97222"))
    assert to_dict(["date"], region_97222_ts.data) == {
        pd.to_datetime("2020-04-01"): {"m2": 10, "location_id": "iso1:us#fips:97222",}
    }
    assert region_97222_ts.latest["m2"] == 11


def test_multi_region_get_counties_and_places():
    ds_in = timeseries.MultiRegionDataset.from_csv(
        io.StringIO(
            "location_id,county,aggregate_level,date,m1,m2\n"
            "iso1:us#fips:97111,Bar County,county,2020-04-02,2,\n"
            "iso1:us#fips:97111,Bar County,county,2020-04-03,3,\n"
            "iso1:us#fips:97222,Foo County,county,2020-04-01,,10\n"
            "iso1:us#fips:9711122,,place,2020-04-02,5,60\n"
            "iso1:us#fips:97,Great State,state,2020-04-01,1,2\n"
            "iso1:us#fips:97111,Bar County,county,,3,\n"
            "iso1:us#fips:9711122,,place,,3,\n"
            "iso1:us#fips:97222,Foo County,county,,,10\n"
            "iso1:us#fips:97,Great State,state,,1,2\n"
        )
    )
    ds_out = ds_in.get_counties_and_places(
        after=pd.to_datetime("2020-04-01")
    ).timeseries.reset_index()
    assert to_dict(["location_id", "date"], ds_out[["location_id", "date", "m1"]]) == {
        ("iso1:us#fips:97111", pd.to_datetime("2020-04-02")): {"m1": 2},
        ("iso1:us#fips:97111", pd.to_datetime("2020-04-03")): {"m1": 3},
        ("iso1:us#fips:9711122", pd.to_datetime("2020-04-02")): {"m1": 5},
    }


def test_multi_region_groupby():
    ts = timeseries.MultiRegionDataset.from_csv(
        io.StringIO(
            "location_id,county,aggregate_level,date,m1,m2\n"
            "iso1:us#fips:97222,Foo County,county,2020-04-01,,10\n"
            "iso1:us#fips:97222,Foo County,county,2020-04-02,,20\n"
            "iso1:us#fips:97,Great State,state,2020-04-01,1,2\n"
            "iso1:us#fips:97222,Foo County,county,,,20\n"
            "iso1:us#fips:97,Great State,state,,1,2\n"
        )
    )

    assert ts.groupby_region()["m2"].last().to_dict() == {
        "iso1:us#fips:97": 2,
        "iso1:us#fips:97222": 20,
    }


def test_one_region_dataset():
    bar_county_row = {
        "location_id": "iso1:us#fips:97111",
        "county": "Bar County",
        "aggregate_level": "county",
        "date": "2020-04-02",
        "m1": 2,
        "m2": pd.NA,
    }
    ts = timeseries.OneRegionTimeseriesDataset(
        Region.from_fips("97111"), pd.DataFrame([bar_county_row]), {}, pd.DataFrame([])
    )
    assert ts.has_one_region() == True

    foo_county_row = {
        "location_id": "iso1:us#fips:97222",
        "county": "Foo County",
        "aggregate_level": "county",
        "date": "2020-04-01",
        "m1": pd.NA,
        "m2": 10,
    }
    with pytest.raises(ValueError):
        timeseries.OneRegionTimeseriesDataset(
            Region.from_fips("97222"),
            pd.DataFrame([bar_county_row, foo_county_row]),
            {},
            pd.DataFrame([]),
        )

    with structlog.testing.capture_logs() as logs:
        ts = timeseries.OneRegionTimeseriesDataset(
            Region.from_fips("97111"),
            pd.DataFrame([], columns="location_id county aggregate_level date m1 m2".split()),
            {},
            pd.DataFrame([]),
        )
    assert [l["event"] for l in logs] == ["Creating OneRegionTimeseriesDataset with zero regions"]
    assert ts.empty


def test_multiregion_provenance():
    m1 = FieldName("m1")
    m2 = FieldName("m2")

    region_97111 = Region.from_fips("97111")
    region_97222 = Region.from_fips("97222")
    region_03 = Region.from_fips("03")
    ds = test_helpers.build_dataset(
        {
            region_97111: {m1: TimeseriesLiteral([1, 2, None], provenance="src11")},
            region_97222: {
                m1: TimeseriesLiteral([None, None, 3], provenance="src21"),
                m2: TimeseriesLiteral([10, None, 30], provenance="src22"),
            },
            region_03: {
                m1: TimeseriesLiteral([None, None, 4], provenance="src31"),
                m2: TimeseriesLiteral([None, None, 40], provenance="src32"),
            },
        },
    )

    # Use loc[...].at[...] as work-around for https://github.com/pandas-dev/pandas/issues/26989
    assert ds.provenance.loc[region_97111.location_id].at["m1"] == "src11"
    assert ds.get_one_region(region_97111).provenance["m1"] == ["src11"]
    assert ds.provenance.loc[region_97222.location_id].at["m2"] == "src22"
    assert ds.get_one_region(region_97222).provenance["m2"] == ["src22"]
    assert ds.provenance.loc[region_03.location_id].at["m2"] == "src32"
    assert ds.get_one_region(region_03).provenance["m2"] == ["src32"]

    counties = ds.get_counties_and_places(after=pd.to_datetime("2020-04-01"))
    assert region_03.location_id not in counties.provenance.index
    assert counties.provenance.loc[region_97222.location_id].at["m1"] == "src21"
    assert counties.get_one_region(region_97222).provenance["m1"] == ["src21"]


def test_one_region_multiple_provenance():
    tag1 = test_helpers.make_tag(date="2020-04-01")
    tag2 = test_helpers.make_tag(date="2020-04-02")
    one_region = test_helpers.build_one_region_dataset(
        {
            CommonFields.ICU_BEDS: TimeseriesLiteral(
                [0, 2, 4], annotation=[tag1, tag2], provenance=["prov1", "prov2"],
            ),
            CommonFields.CASES: [100, 200, 300],
        }
    )

    assert set(one_region.annotations_all_bucket(CommonFields.ICU_BEDS)) == {tag1, tag2}
    assert sorted(one_region.provenance[CommonFields.ICU_BEDS]) == ["prov1", "prov2"]


def test_add_aggregate_level():
    ts_df = read_csv_and_index_fips_date("fips,date,m1,m2\n" "36061,2020-04-02,2,\n").reset_index()
    multiregion = timeseries.MultiRegionDataset.from_fips_timeseries_df(ts_df)
    assert multiregion.geo_data.aggregate_level.to_list() == ["county"]


def test_fips_not_in_geo_data_csv_raises():
    df = test_helpers.read_csv_str(
        "       location_id,       date,       cases\n"
        "iso1:us#fips:06010, 2020-04-01,         100\n",
        skip_spaces=True,
    )

    with pytest.raises(AssertionError):
        timeseries.MultiRegionDataset.from_timeseries_df(df)


def test_append_regions():
    ts_fips = timeseries.MultiRegionDataset.from_csv(
        io.StringIO(
            "location_id,date,county,aggregate_level,m1,m2\n"
            "iso1:us#fips:97111,2020-04-02,Bar County,county,2,\n"
            "iso1:us#fips:97111,2020-04-03,Bar County,county,3,\n"
            "iso1:us#fips:97222,2020-04-04,Foo County,county,,11\n"
            "iso1:us#fips:97111,,Bar County,county,3,\n"
            "iso1:us#fips:97222,,Foo County,county,,11\n"
        )
    ).add_provenance_csv(
        io.StringIO("location_id,variable,provenance\n" "iso1:us#fips:97111,m1,prov97111m1\n")
    )
    ts_cbsa = timeseries.MultiRegionDataset.from_csv(
        io.StringIO(
            "location_id,date,m2\n"
            "iso1:us#cbsa:10100,2020-04-02,2\n"
            "iso1:us#cbsa:10100,2020-04-03,3\n"
            "iso1:us#cbsa:20300,2020-04-03,4\n"
            "iso1:us#cbsa:10100,,3\n"
            "iso1:us#cbsa:20300,,4\n"
        )
    ).add_provenance_csv(
        io.StringIO("location_id,variable,provenance\n" "iso1:us#cbsa:20300,m1,prov20200m2\n")
    )
    # Check that merge is symmetric
    ts_merged_1 = ts_fips.append_regions(ts_cbsa)
    ts_merged_2 = ts_cbsa.append_regions(ts_fips)
    test_helpers.assert_dataset_like(ts_merged_1, ts_merged_2)

    ts_expected = timeseries.MultiRegionDataset.from_csv(
        io.StringIO(
            "location_id,date,county,aggregate_level,m1,m2\n"
            "iso1:us#cbsa:10100,2020-04-02,,,,2\n"
            "iso1:us#cbsa:10100,2020-04-03,,,,3\n"
            "iso1:us#cbsa:20300,2020-04-03,,,,4\n"
            "iso1:us#cbsa:10100,,,,,3\n"
            "iso1:us#cbsa:20300,,,,,4\n"
            "iso1:us#fips:97111,2020-04-02,Bar County,county,2,\n"
            "iso1:us#fips:97111,2020-04-03,Bar County,county,3,\n"
            "iso1:us#fips:97222,2020-04-04,Foo County,county,,11\n"
            "iso1:us#fips:97111,,Bar County,county,3,\n"
            "iso1:us#fips:97222,,Foo County,county,,11\n"
        )
    ).add_provenance_csv(
        io.StringIO(
            "location_id,variable,provenance\n"
            "iso1:us#fips:97111,m1,prov97111m1\n"
            "iso1:us#cbsa:20300,m1,prov20200m2\n"
        )
    )
    test_helpers.assert_dataset_like(ts_merged_1, ts_expected)


def test_append_regions_with_buckets():
    region_cbsa = Region.from_cbsa_code("10100")
    region_la = Region.from_fips("06037")
    region_sf = Region.from_fips("06075")
    m1 = FieldName("m1")
    m2 = FieldName("m2")
    age_40s = DemographicBucket("age:40-49")
    data_county = {
        region_la: {
            m1: {
                age_40s: TimeseriesLiteral([1, 2], annotation=[test_helpers.make_tag()]),
                DemographicBucket.ALL: [2, 3],
            }
        },
        region_sf: {m1: [3, 4]},
    }
    data_cbsa = {region_cbsa: {m2: [5, 6]}}
    ds_county = test_helpers.build_dataset(data_county)
    ds_cbsa = test_helpers.build_dataset(data_cbsa)

    ds_out_1 = ds_county.append_regions(ds_cbsa)
    ds_out_2 = ds_cbsa.append_regions(ds_county)

    ds_expected = test_helpers.build_dataset({**data_cbsa, **data_county})

    test_helpers.assert_dataset_like(ds_out_1, ds_expected)
    test_helpers.assert_dataset_like(ds_out_2, ds_expected)


def test_append_regions_duplicate_region_raises():
    ts1 = timeseries.MultiRegionDataset.from_csv(
        io.StringIO(
            "location_id,date,county,aggregate_level,m1,m2\n"
            "iso1:us#fips:97111,2020-04-02,Bar County,county,2,\n"
        )
    )
    ts2 = timeseries.MultiRegionDataset.from_csv(
        io.StringIO(
            "location_id,date,county,aggregate_level,m1,m2\n"
            "iso1:us#fips:97111,2020-04-03,Bar County,county,2,\n"
        )
    )
    with pytest.raises(ValueError):
        ts1.append_regions(ts2)


def test_timeseries_long():
    """Test timeseries_long where all data has bucket `all`"""
    region_cbsa = Region.from_cbsa_code("10100")
    region_county = Region.from_fips("97111")
    ds = test_helpers.build_dataset(
        {
            region_county: {FieldName("m1"): [2, None, 4]},
            region_cbsa: {FieldName("m2"): [2, 3, None]},
        },
        start_date="2020-04-02",
    )

    expected = test_helpers.read_csv_str(
        "       location_id,      date,variable,value\n"
        "iso1:us#cbsa:10100,2020-04-02,      m2,    2\n"
        "iso1:us#cbsa:10100,2020-04-03,      m2,    3\n"
        "iso1:us#fips:97111,2020-04-02,      m1,    2\n"
        "iso1:us#fips:97111,2020-04-04,      m1,    4\n",
        skip_spaces=True,
        dtype={"value": float},
    )
    long_series = ds.timeseries_bucketed_long
    assert long_series.index.names == [
        CommonFields.LOCATION_ID,
        PdFields.DEMOGRAPHIC_BUCKET,
        CommonFields.DATE,
        PdFields.VARIABLE,
    ]
    assert long_series.name == PdFields.VALUE
    long_df = long_series.xs("all", level=PdFields.DEMOGRAPHIC_BUCKET).reset_index()
    pd.testing.assert_frame_equal(long_df, expected, check_like=True)


def test_timeseries_bucketed_long():
    region_cbsa = Region.from_cbsa_code("10100")
    region_county = Region.from_fips("97111")
    bucket_age_0 = DemographicBucket("age:0-9")
    bucket_age_10 = DemographicBucket("age:10-19")
    bucket_all = DemographicBucket("all")
    ds = test_helpers.build_dataset(
        {
            region_county: {
                FieldName("m1"): {
                    bucket_age_0: [4, 5, 6],
                    bucket_age_10: [None, None, 7],
                    bucket_all: [2, None, 4],
                }
            },
            region_cbsa: {FieldName("m2"): [2, 3, None]},
        },
        start_date="2020-04-02",
    )

    expected = test_helpers.read_csv_str(
        "       location_id,demographic_bucket,      date,variable,value\n"
        "iso1:us#cbsa:10100,               all,2020-04-02,      m2,    2\n"
        "iso1:us#cbsa:10100,               all,2020-04-03,      m2,    3\n"
        "iso1:us#fips:97111,           age:0-9,2020-04-02,      m1,    4\n"
        "iso1:us#fips:97111,           age:0-9,2020-04-03,      m1,    5\n"
        "iso1:us#fips:97111,           age:0-9,2020-04-04,      m1,    6\n"
        "iso1:us#fips:97111,         age:10-19,2020-04-04,      m1,    7\n"
        "iso1:us#fips:97111,               all,2020-04-02,      m1,    2\n"
        "iso1:us#fips:97111,               all,2020-04-04,      m1,    4\n",
        skip_spaces=True,
        dtype={"value": float},
    )
    long_series = ds.timeseries_bucketed_long
    assert long_series.index.names == [
        CommonFields.LOCATION_ID,
        PdFields.DEMOGRAPHIC_BUCKET,
        CommonFields.DATE,
        PdFields.VARIABLE,
    ]
    assert long_series.name == PdFields.VALUE
    pd.testing.assert_frame_equal(long_series.reset_index(), expected, check_like=True)


def test_timeseries_distribution_long():
    bucket_age_0 = DemographicBucket("age:0-9")
    bucket_age_10 = DemographicBucket("age:10-19")
    bucket_all = DemographicBucket("all")
    bucket_blueman = DemographicBucket("color;gender:blue;man")
    ds = test_helpers.build_default_region_dataset(
        {
            FieldName("m1"): {
                bucket_age_0: [1, 2, 3],
                bucket_age_10: [None, None, 4],
                bucket_all: [5, None, 6],
                bucket_blueman: [7, None, None],
            }
        }
    )

    long_series = ds.timeseries_distribution_long
    assert long_series.name == PdFields.VALUE
    assert long_series.index.names == [
        CommonFields.LOCATION_ID,
        PdFields.DEMOGRAPHIC_BUCKET,
        CommonFields.DATE,
        PdFields.DISTRIBUTION,
        PdFields.VARIABLE,
    ]
    expected = test_helpers.read_csv_str(
        "       location_id,   demographic_bucket,      date,distribution,variable,value\n"
        "iso1:us#fips:97222,              age:0-9,2020-04-01,         age,      m1,    1\n"
        "iso1:us#fips:97222,              age:0-9,2020-04-02,         age,      m1,    2\n"
        "iso1:us#fips:97222,              age:0-9,2020-04-03,         age,      m1,    3\n"
        "iso1:us#fips:97222,            age:10-19,2020-04-03,         age,      m1,    4\n"
        "iso1:us#fips:97222,                  all,2020-04-01,         all,      m1,    5\n"
        "iso1:us#fips:97222,                  all,2020-04-03,         all,      m1,    6\n"
        "iso1:us#fips:97222,color;gender:blue;man,2020-04-01,color;gender,      m1,    7\n",
        skip_spaces=True,
        dtype={"value": float},
    )
    pd.testing.assert_frame_equal(long_series.reset_index(), expected, check_like=True)


def test_timeseries_wide_dates():
    region_cbsa = Region.from_cbsa_code("10100")
    region_fips = Region.from_fips("97111")
    m1 = FieldName("m1")
    m2 = FieldName("m2")
    ds = test_helpers.build_dataset(
        {region_cbsa: {m2: [2, 3]}, region_fips: {m1: [2, None, 4]}},
        static_by_region_then_field_name={region_fips: {CommonFields.COUNTY: "Bar County", m1: 4}},
        start_date="2020-04-02",
    )

    # TODO(tom): Delete this test of _timeseries_not_bucketed_wide_dates which is no longer
    #  accessed from outside timeseries when there are other tests for from_timeseries_wide_dates_df
    ds_wide = ds._timeseries_not_bucketed_wide_dates
    assert ds_wide.index.names == [CommonFields.LOCATION_ID, PdFields.VARIABLE]
    assert ds_wide.columns.names == [CommonFields.DATE]

    expected = (
        pd.read_csv(
            io.StringIO(
                "location_id,variable,2020-04-02,2020-04-03,2020-04-04\n"
                "iso1:us#cbsa:10100,m2,2,3,\n"
                "iso1:us#fips:97111,m1,2,,4\n"
            ),
        )
        .set_index(ds_wide.index.names)
        .rename_axis(columns="date")
        .astype(float)
    )
    expected.columns = pd.to_datetime(expected.columns)

    pd.testing.assert_frame_equal(ds_wide, expected)

    # Recreate the dataset using `from_timeseries_wide_dates_df`.
    ds_recreated = timeseries.MultiRegionDataset.from_timeseries_wide_dates_df(
        ds_wide
    ).add_static_values(ds.static.reset_index())
    test_helpers.assert_dataset_like(ds, ds_recreated)

    assert ds.get_timeseries_not_bucketed_wide_dates(m1).loc[region_fips.location_id, :].replace(
        {np.nan: None}
    ).to_list() == [2, None, 4]
    assert ds.get_timeseries_bucketed_wide_dates(m1).loc[
        (region_fips.location_id, DemographicBucket.ALL), :
    ].replace({np.nan: None}).to_list() == [2, None, 4]


def test_timeseries_wide_dates_empty():
    m1 = FieldName("m1")
    ds = timeseries.MultiRegionDataset.from_csv(
        io.StringIO(
            "location_id,date,county,aggregate_level,m1,m2\n"
            "iso1:us#cbsa:10100,,,,,3\n"
            "iso1:us#fips:97111,,Bar County,county,4,\n"
        )
    )

    assert ds.get_timeseries_not_bucketed_wide_dates(m1).empty
    assert ds.get_timeseries_bucketed_wide_dates(m1).empty
    assert ds.get_timeseries_not_bucketed_wide_dates(CommonFields.CASES).empty
    assert ds.get_timeseries_bucketed_wide_dates(CommonFields.CASES).empty


def test_write_read_wide_dates_csv_compare_literal(tmpdir):
    pointer = _make_dataset_pointer(tmpdir)

    region_as = Region.from_state("AS")
    region_sf = Region.from_fips("06075")
    metrics_as = {
        CommonFields.ICU_BEDS: TimeseriesLiteral([0, 2, 4], provenance="pt_src1"),
        CommonFields.CASES: [100, 200, 300],
    }
    metrics_sf = {
        CommonFields.DEATHS: TimeseriesLiteral([1, 2, None], provenance="pt_src2"),
        CommonFields.CASES: [None, 210, 310],
    }
    dataset_in = test_helpers.build_dataset({region_as: metrics_as, region_sf: metrics_sf})

    dataset_in.write_to_dataset_pointer(pointer)

    # Compare written file with a string literal so a test fails if something changes in how the
    # file is written. The literal contains spaces to align the columns in the source.
    assert pointer.path_wide_dates().read_text() == (
        "     location_id,variable,demographic_bucket,provenance,2020-04-03,2020-04-02,2020-04-01\n"
        "iso1:us#iso2:us-as, cases,               all,          ,       300,       200,       100\n"
        "iso1:us#iso2:us-as,icu_beds,             all,   pt_src1,         4,         2,         0\n"
        "iso1:us#iso2:us-ca#fips:06075,cases,     all,          ,       310,       210,\n"
        "iso1:us#iso2:us-ca#fips:06075,deaths,    all,   pt_src2,          ,         2,         1\n"
    ).replace(" ", "")

    dataset_read = timeseries.MultiRegionDataset.read_from_pointer(pointer)
    test_helpers.assert_dataset_like(dataset_read, dataset_in)

    # Check that a file without the demographic_bucket column (as written before
    # https://github.com/covid-projections/covid-data-model/pull/1021) can be read.
    pointer.path_wide_dates().write_text(
        "                  location_id,variable,provenance,2020-04-03,2020-04-02,2020-04-01\n"
        "           iso1:us#iso2:us-as,   cases,          ,       300,       200,       100\n"
        "           iso1:us#iso2:us-as,icu_beds,   pt_src1,         4,         2,         0\n"
        "iso1:us#iso2:us-ca#fips:06075,   cases,          ,       310,       210,\n"
        "iso1:us#iso2:us-ca#fips:06075,  deaths,   pt_src2,          ,         2,         1\n".replace(
            " ", ""
        )
    )
    dataset_without_bucket_read = timeseries.MultiRegionDataset.read_from_pointer(pointer)
    test_helpers.assert_dataset_like(dataset_without_bucket_read, dataset_in)


def test_write_read_wide_dates_csv_with_annotation(tmpdir):
    pointer = _make_dataset_pointer(tmpdir)

    region = Region.from_state("AS")
    metrics = {
        CommonFields.ICU_BEDS: TimeseriesLiteral(
            [0, 2, 4],
            annotation=[
                test_helpers.make_tag(date="2020-04-01"),
                test_helpers.make_tag(TagType.ZSCORE_OUTLIER, date="2020-04-02"),
            ],
        ),
        CommonFields.CASES: [100, 200, 300],
    }
    dataset_in = test_helpers.build_dataset({region: metrics})

    dataset_in.write_to_dataset_pointer(pointer)
    dataset_read = timeseries.MultiRegionDataset.read_from_pointer(pointer)

    test_helpers.assert_dataset_like(dataset_read, dataset_in)


def test_write_read_dataset_pointer_with_provenance_list(tmpdir):
    pointer = _make_dataset_pointer(tmpdir)

    dataset_in = test_helpers.build_default_region_dataset(
        {
            CommonFields.ICU_BEDS: TimeseriesLiteral(
                [0, 2, 4],
                annotation=[
                    test_helpers.make_tag(date="2020-04-01"),
                    test_helpers.make_tag(date="2020-04-02"),
                ],
                provenance=["prov1", "prov2"],
            ),
            CommonFields.CASES: [100, 200, 300],
        }
    )

    dataset_in.write_to_dataset_pointer(pointer)
    dataset_read = timeseries.MultiRegionDataset.read_from_pointer(pointer)

    test_helpers.assert_dataset_like(dataset_read, dataset_in)


def test_write_read_wide_with_buckets(tmpdir):
    pointer = _make_dataset_pointer(tmpdir)

    all_bucket = DemographicBucket("all")
    age_20s = DemographicBucket("age:20-29")
    age_30s = DemographicBucket("age:30-39")
    region_as = Region.from_state("AS")
    region_sf = Region.from_fips("06075")
    metrics_as = {
        CommonFields.ICU_BEDS: TimeseriesLiteral(
            [0, 2, 4],
            annotation=[
                test_helpers.make_tag(date="2020-04-01"),
                test_helpers.make_tag(TagType.ZSCORE_OUTLIER, date="2020-04-02"),
            ],
        ),
        CommonFields.CASES: [100, 200, 300],
    }
    metrics_sf = {
        CommonFields.CASES: {
            age_20s: TimeseriesLiteral([3, 4, 5], source=taglib.Source(type="MySource")),
            age_30s: [4, 5, 6],
            all_bucket: [1, 2, 3],
        }
    }
    dataset_in = test_helpers.build_dataset({region_as: metrics_as, region_sf: metrics_sf})

    dataset_in.write_to_dataset_pointer(pointer)
    dataset_read = timeseries.MultiRegionDataset.read_from_pointer(pointer)

    test_helpers.assert_dataset_like(dataset_read, dataset_in)


def test_timeseries_drop_stale_timeseries_entire_region():
    ds_in = timeseries.MultiRegionDataset.from_csv(
        io.StringIO(
            "location_id,date,county,aggregate_level,m1,m2\n"
            "iso1:us#cbsa:10100,2020-04-02,,,,2\n"
            "iso1:us#cbsa:10100,2020-04-03,,,,3\n"
            "iso1:us#cbsa:10100,,,,,3\n"
            "iso1:us#fips:97111,2020-04-02,Bar County,county,2,\n"
            "iso1:us#fips:97111,2020-04-04,Bar County,county,4,\n"
            "iso1:us#fips:97111,,Bar County,county,4,\n"
        )
    )

    ds_out = ds_in.drop_stale_timeseries(pd.to_datetime("2020-04-04"))

    ds_expected = timeseries.MultiRegionDataset.from_csv(
        io.StringIO(
            "location_id,date,county,aggregate_level,m1,m2\n"
            "iso1:us#cbsa:10100,,,,,3\n"
            "iso1:us#fips:97111,2020-04-02,Bar County,county,2,\n"
            "iso1:us#fips:97111,2020-04-04,Bar County,county,4,\n"
            "iso1:us#fips:97111,,Bar County,county,4,\n"
        )
    )
    test_helpers.assert_dataset_like(ds_out, ds_expected)


def test_timeseries_drop_stale_timeseries_one_metric():
    csv_in = (
        "location_id,date,county,aggregate_level,m1,m2\n"
        "iso1:us#cbsa:10100,2020-04-02,,,11,2\n"
        "iso1:us#cbsa:10100,2020-04-03,,,,3\n"
        "iso1:us#cbsa:10100,,,,,3\n"
        "iso1:us#fips:97111,2020-04-02,Bar County,county,2,\n"
        "iso1:us#fips:97111,2020-04-04,Bar County,county,4,\n"
        "iso1:us#fips:97111,,Bar County,county,4,\n"
    )
    ds_in = timeseries.MultiRegionDataset.from_csv(io.StringIO(csv_in)).add_provenance_csv(
        io.StringIO(
            "location_id,variable,provenance\n"
            "iso1:us#cbsa:10100,m1,m1-10100prov\n"
            "iso1:us#cbsa:10100,m2,m2-10100prov\n"
            "iso1:us#fips:97111,m1,m1-97111prov\n"
        )
    )

    ds_out = ds_in.drop_stale_timeseries(pd.to_datetime("2020-04-03"))

    # The only timeseries that is stale with cutoff of 4/3 is the CBSA m1. The expected
    # dataset is the same as the input with "11" removed from the timeseries and
    # corresponding provenance removed.
    ds_expected = timeseries.MultiRegionDataset.from_csv(
        io.StringIO(csv_in.replace(",11,", ",,"))
    ).add_provenance_csv(
        io.StringIO(
            "location_id,variable,provenance\n"
            "iso1:us#cbsa:10100,m2,m2-10100prov\n"
            "iso1:us#fips:97111,m1,m1-97111prov\n"
        )
    )
    test_helpers.assert_dataset_like(ds_out, ds_expected)


def test_timeseries_drop_stale_timeseries_with_tag():
    region = Region.from_state("TX")
    values_recent = [100, 200, 300, 400]
    values_stale = [100, 200, None, None]
    ts_recent = TimeseriesLiteral(values_recent, annotation=[test_helpers.make_tag()])
    ts_stale = TimeseriesLiteral(values_stale, annotation=[test_helpers.make_tag()])

    dataset_in = test_helpers.build_dataset(
        {region: {CommonFields.CASES: ts_recent, CommonFields.DEATHS: ts_stale}}
    )

    dataset_out = dataset_in.drop_stale_timeseries(pd.to_datetime("2020-04-03"))

    assert len(dataset_out.tag) == 1
    # drop_stale_timeseries preserves the empty DEATHS column so add it to dataset_expected
    dataset_expected = test_helpers.build_dataset(
        {region: {CommonFields.CASES: ts_recent}}, timeseries_columns=[CommonFields.DEATHS]
    )
    test_helpers.assert_dataset_like(dataset_out, dataset_expected)


def test_append_region_and_get_regions_subset_with_tag():
    region_tx = Region.from_state("TX")
    region_sf = Region.from_fips("06075")
    values = [100, 200, 300, 400]
    ts_with_tag = TimeseriesLiteral(values, annotation=[test_helpers.make_tag()])

    dataset_tx = test_helpers.build_dataset({region_tx: {CommonFields.CASES: ts_with_tag}})
    dataset_sf = test_helpers.build_dataset({region_sf: {CommonFields.CASES: ts_with_tag}})

    dataset_appended = dataset_tx.append_regions(dataset_sf)

    assert len(dataset_appended.tag) == 2
    dataset_tx_and_sf = test_helpers.build_dataset(
        {region_tx: {CommonFields.CASES: ts_with_tag}, region_sf: {CommonFields.CASES: ts_with_tag}}
    )
    test_helpers.assert_dataset_like(dataset_appended, dataset_tx_and_sf)

    dataset_out = dataset_tx_and_sf.get_regions_subset([region_tx])
    assert len(dataset_out.tag) == 1
    test_helpers.assert_dataset_like(dataset_out, dataset_tx)


def test_one_region_annotations():
    region_tx = Region.from_state("TX")
    region_sf = Region.from_fips("06075")
    values = [100, 200, 300, 400]
    tag1 = test_helpers.make_tag(date="2020-04-01")
    tag2a = test_helpers.make_tag(date="2020-04-02")
    tag2b = test_helpers.make_tag(date="2020-04-03")

    dataset_tx_and_sf = test_helpers.build_dataset(
        {
            region_tx: {CommonFields.CASES: (TimeseriesLiteral(values, annotation=[tag1]))},
            region_sf: {CommonFields.CASES: (TimeseriesLiteral(values, annotation=[tag2a, tag2b]))},
        }
    )

    # get_one_region and iter_one_regions use separate code to split up the tags. Test both of them.
    one_region_tx = dataset_tx_and_sf.get_one_region(region_tx)
    assert one_region_tx.annotations_all_bucket(CommonFields.CASES) == [tag1]
    one_region_sf = dataset_tx_and_sf.get_one_region(region_sf)
    assert one_region_sf.annotations_all_bucket(CommonFields.CASES) == [
        tag2a,
        tag2b,
    ]
    assert set(one_region_sf.sources_all_bucket(CommonFields.CASES)) == set()

    assert {
        region: one_region_dataset.annotations_all_bucket(CommonFields.CASES)
        for region, one_region_dataset in dataset_tx_and_sf.iter_one_regions()
    } == {region_sf: [tag2a, tag2b], region_tx: [tag1],}


def test_one_region_empty_annotations():
    one_region = test_helpers.build_one_region_dataset({CommonFields.CASES: [100, 200, 300]})

    assert one_region.annotations_all_bucket(CommonFields.CASES) == []
    assert one_region.source_url == {}
    assert one_region.provenance == {}
    assert set(one_region.sources_all_bucket(CommonFields.ICU_BEDS)) == set()
    assert set(one_region.sources_all_bucket(CommonFields.CASES)) == set()


def test_one_region_tag_objects_series():
    values = [100, 200]
    tag1 = test_helpers.make_tag(TagType.ZSCORE_OUTLIER, date="2020-04-01")
    tag2a = test_helpers.make_tag(date="2020-04-02")
    tag2b = test_helpers.make_tag(date="2020-04-03")

    one_region = test_helpers.build_one_region_dataset(
        {
            CommonFields.CASES: TimeseriesLiteral(values, annotation=[tag1]),
            CommonFields.ICU_BEDS: TimeseriesLiteral(values, provenance="prov1"),
            CommonFields.DEATHS: TimeseriesLiteral(values, annotation=[tag2a, tag2b]),
        }
    )

    assert isinstance(one_region.tag_objects_series, pd.Series)
    assert one_region.tag.index.equals(one_region.tag_objects_series.index)
    assert set(one_region.tag_objects_series.reset_index().itertuples(index=False)) == {
        (CommonFields.CASES, DemographicBucket.ALL, tag1.tag_type, tag1),
        (
            CommonFields.ICU_BEDS,
            DemographicBucket.ALL,
            "provenance",
            taglib.ProvenanceTag(source="prov1"),
        ),
        (CommonFields.DEATHS, DemographicBucket.ALL, tag2a.tag_type, tag2a),
        (CommonFields.DEATHS, DemographicBucket.ALL, tag2b.tag_type, tag2b),
    }


def test_one_region_tag_objects_series_empty():
    one_region = test_helpers.build_one_region_dataset({CommonFields.CASES: [1, 2, 3]})
    assert one_region.tag.empty
    assert isinstance(one_region.tag_objects_series, pd.Series)
    assert one_region.tag_objects_series.empty


def test_timeseries_tag_objects_series():
    values = [100, 200]
    tag1 = test_helpers.make_tag(TagType.ZSCORE_OUTLIER, date="2020-04-01")
    tag2a = test_helpers.make_tag(date="2020-04-02")
    tag2b = test_helpers.make_tag(date="2020-04-03")
    url_str = UrlStr("http://foo.com/1")
    source_obj = taglib.Source("source_with_url", url=url_str)

    ds = test_helpers.build_default_region_dataset(
        {
            CommonFields.CASES: TimeseriesLiteral(values, annotation=[tag1]),
            CommonFields.ICU_BEDS: TimeseriesLiteral(values, source=source_obj),
            CommonFields.DEATHS: TimeseriesLiteral(values, annotation=[tag2a, tag2b]),
            CommonFields.TOTAL_TESTS: values,
        }
    )

    assert isinstance(ds.tag_objects_series, pd.Series)
    assert ds.tag.index.equals(ds.tag_objects_series.index)
    location_id = test_helpers.DEFAULT_REGION.location_id
    assert set(ds.tag_objects_series.reset_index().itertuples(index=False)) == {
        (location_id, CommonFields.CASES, DemographicBucket.ALL, tag1.tag_type, tag1),
        (location_id, CommonFields.ICU_BEDS, DemographicBucket.ALL, TagType.SOURCE, source_obj),
        (location_id, CommonFields.DEATHS, DemographicBucket.ALL, tag2a.tag_type, tag2a),
        (location_id, CommonFields.DEATHS, DemographicBucket.ALL, tag2b.tag_type, tag2b),
    }


def test_timeseries_latest_values():
    dataset = timeseries.MultiRegionDataset.from_csv(
        io.StringIO(
            "location_id,date,county,aggregate_level,m1,m2\n"
            "iso1:us#cbsa:10100,2020-04-02,,,,2\n"
            "iso1:us#cbsa:10100,2020-04-03,,,10,3\n"
            "iso1:us#cbsa:10100,2020-04-04,,,,1\n"
            "iso1:us#cbsa:10100,,,,,4\n"
            "iso1:us#fips:97111,2020-04-02,Bar County,county,2,\n"
            "iso1:us#fips:97111,2020-04-04,Bar County,county,4,\n"
            "iso1:us#fips:97111,,Bar County,county,5,\n"
        )
    )

    # Check bulk access via _timeseries_latest_values
    expected = pd.read_csv(
        io.StringIO("location_id,m1,m2\n" "iso1:us#cbsa:10100,10,1\n" "iso1:us#fips:97111,4,\n")
    )
    latest_from_timeseries = dataset._timeseries_latest_values().reset_index()
    pd.testing.assert_frame_equal(
        latest_from_timeseries, expected, check_like=True, check_dtype=False
    )

    # Check access to timeseries latests values via get_one_region
    region_10100 = dataset.get_one_region(Region.from_cbsa_code("10100"))
    assert region_10100.latest == {
        "aggregate_level": "cbsa",
        "county": None,
        "country": "USA",
        "fips": "10100",
        "state": None,
        "m1": 10,  # Derived from timeseries
        "m2": 4,  # Explicitly in recent values
    }
    region_97111 = dataset.get_one_region(Region.from_fips("97111"))
    assert region_97111.latest == {
        "aggregate_level": "county",
        "county": "Bar County",
        "country": "USA",
        "fips": "97111",
        "state": "ZZ",
        "m1": 5,
        "m2": None,
    }


def test_timeseries_latest_values_copied_to_static():
    dataset = timeseries.MultiRegionDataset.from_csv(
        io.StringIO(
            "location_id,date,county,aggregate_level,t1,s1\n"
            "iso1:us#cbsa:10100,2020-04-02,,,,2\n"
            "iso1:us#cbsa:10100,2020-04-03,,,10,3\n"
            "iso1:us#cbsa:10100,2020-04-04,,,,1\n"
            "iso1:us#cbsa:10100,,,,,4\n"
            "iso1:us#fips:97111,2020-04-02,Bar County,county,2,\n"
            "iso1:us#fips:97111,2020-04-04,Bar County,county,4,\n"
            "iso1:us#fips:97111,,Bar County,county,,\n"
        )
    )

    # Check access to latest values as copied to static
    t1 = FieldName("t1")
    s1 = FieldName("s1")
    dataset_t1_latest_in_static = dataset.latest_in_static(t1)
    assert dataset_t1_latest_in_static.static.loc["iso1:us#cbsa:10100", t1] == 10
    assert dataset_t1_latest_in_static.static.loc["iso1:us#fips:97111", t1] == 4

    # Trying to copy the latest values of s1 fails because s1 already has a real value in static.
    # See also longer comment where the ValueError is raised.
    with pytest.raises(ValueError):
        dataset.latest_in_static(s1)


def test_join_columns():
    ts_1 = timeseries.MultiRegionDataset.from_csv(
        io.StringIO(
            "location_id,date,county,aggregate_level,m1\n"
            "iso1:us#cbsa:10100,2020-04-02,,,\n"
            "iso1:us#cbsa:10100,2020-04-03,,,\n"
            "iso1:us#cbsa:10100,,,,\n"
            "iso1:us#fips:97111,2020-04-02,Bar County,county,2\n"
            "iso1:us#fips:97111,2020-04-04,Bar County,county,4\n"
            "iso1:us#fips:97111,,Bar County,county,4\n"
        )
    ).add_provenance_csv(
        io.StringIO(
            "location_id,variable,provenance\n"
            "iso1:us#cbsa:10100,m1,ts110100prov\n"
            "iso1:us#fips:97111,m1,ts197111prov\n"
        )
    )
    ts_2 = timeseries.MultiRegionDataset.from_csv(
        io.StringIO(
            "location_id,date,county,aggregate_level,m2\n"
            "iso1:us#cbsa:10100,2020-04-02,,,2\n"
            "iso1:us#cbsa:10100,2020-04-03,,,3\n"
            "iso1:us#fips:97111,2020-04-02,Bar County,county,\n"
            "iso1:us#fips:97111,2020-04-04,Bar County,county,\n"
        )
    ).add_provenance_csv(
        io.StringIO(
            "location_id,variable,provenance\n"
            "iso1:us#cbsa:10100,m2,ts110100prov\n"
            "iso1:us#fips:97111,m2,ts197111prov\n"
        )
    )
    ts_expected = timeseries.MultiRegionDataset.from_csv(
        io.StringIO(
            "location_id,date,county,aggregate_level,m1,m2\n"
            "iso1:us#cbsa:10100,2020-04-02,,,,2\n"
            "iso1:us#cbsa:10100,2020-04-03,,,,3\n"
            "iso1:us#cbsa:10100,,,,,\n"
            "iso1:us#fips:97111,2020-04-02,Bar County,county,2,\n"
            "iso1:us#fips:97111,2020-04-04,Bar County,county,4,\n"
            "iso1:us#fips:97111,,Bar County,county,4,\n"
        )
    ).add_provenance_csv(
        io.StringIO(
            "location_id,variable,provenance\n"
            "iso1:us#cbsa:10100,m1,ts110100prov\n"
            "iso1:us#cbsa:10100,m2,ts110100prov\n"
            "iso1:us#fips:97111,m1,ts197111prov\n"
            "iso1:us#fips:97111,m2,ts197111prov\n"
        )
    )
    ts_joined = ts_1.join_columns(ts_2)
    test_helpers.assert_dataset_like(ts_joined, ts_expected, drop_na_latest=True)

    ts_joined = ts_2.join_columns(ts_1)
    test_helpers.assert_dataset_like(ts_joined, ts_expected, drop_na_latest=True)

    with pytest.raises(ValueError):
        # Raises because the same column is in both datasets
        ts_2.join_columns(ts_2)

    # geo attributes, such as aggregation level and county name, generally appear in geo-data.csv
    # instead of MultiRegionDataset so they don't need special handling in join_columns.


def test_join_columns_missing_regions():
    ts_1 = timeseries.MultiRegionDataset.from_csv(
        io.StringIO(
            "location_id,date,county,aggregate_level,m1\n"
            "iso1:us#cbsa:10100,2020-04-02,,,\n"
            "iso1:us#cbsa:10100,2020-04-03,,,\n"
            "iso1:us#cbsa:10100,,,,\n"
            "iso1:us#fips:97111,2020-04-02,Bar County,county,2\n"
            "iso1:us#fips:97111,2020-04-04,Bar County,county,4\n"
            "iso1:us#fips:97111,,Bar County,county,4\n"
        )
    )
    ts_2 = timeseries.MultiRegionDataset.from_csv(
        io.StringIO(
            "location_id,date,county,aggregate_level,m2\n" "iso1:us#cbsa:10100,2020-04-02,,,2\n"
        )
    )
    ts_expected = timeseries.MultiRegionDataset.from_csv(
        io.StringIO(
            "location_id,date,county,aggregate_level,m1,m2\n"
            "iso1:us#cbsa:10100,2020-04-02,,,,2\n"
            "iso1:us#cbsa:10100,2020-04-03,,,,\n"
            "iso1:us#cbsa:10100,,,,,\n"
            "iso1:us#fips:97111,2020-04-02,Bar County,county,2,\n"
            "iso1:us#fips:97111,2020-04-04,Bar County,county,4,\n"
            "iso1:us#fips:97111,,Bar County,county,4,\n"
        )
    )
    ts_joined = ts_1.join_columns(ts_2)
    test_helpers.assert_dataset_like(ts_joined, ts_expected, drop_na_latest=True)


def test_join_columns_with_buckets():
    m1 = FieldName("m1")
    m2 = FieldName("m2")
    age20s = DemographicBucket("age:20-29")

    m1_data = {m1: {age20s: [1, 2, 3]}}
    ds_1 = test_helpers.build_default_region_dataset(m1_data)
    m2_data = {m2: {age20s: [4, 5, 6], DemographicBucket.ALL: [7, 8, 9]}}
    ds_2 = test_helpers.build_default_region_dataset(m2_data)

    with pytest.raises(ValueError):
        ds_1.join_columns(ds_1)

    ds_expected = test_helpers.build_default_region_dataset({**m1_data, **m2_data})

    ds_joined = ds_1.join_columns(ds_2)
    test_helpers.assert_dataset_like(ds_joined, ds_expected)


def test_join_columns_with_static():
    m1 = FieldName("m1")
    m2 = FieldName("m2")

    ds_1 = test_helpers.build_default_region_dataset({}, static={m1: 1})
    ds_2 = test_helpers.build_default_region_dataset({}, static={m2: 2})

    with pytest.raises(ValueError):
        ds_1.join_columns(ds_1)

    ds_expected = test_helpers.build_default_region_dataset({}, static={m1: 1, m2: 2})

    ds_joined = ds_1.join_columns(ds_2)
    test_helpers.assert_dataset_like(ds_joined, ds_expected)

    ds_joined = ds_2.join_columns(ds_1)
    test_helpers.assert_dataset_like(ds_joined, ds_expected)


def test_iter_one_region():
    ts = timeseries.MultiRegionDataset.from_csv(
        io.StringIO(
            "location_id,date,county,aggregate_level,m1\n"
            "iso1:us#cbsa:10100,2020-04-02,,,\n"
            "iso1:us#cbsa:10100,2020-04-03,,,\n"
            "iso1:us#cbsa:10100,,,,\n"
            "iso1:us#fips:97111,2020-04-02,Bar County,county,2\n"
            "iso1:us#fips:97111,2020-04-04,Bar County,county,4\n"
            "iso1:us#fips:97111,,Bar County,county,4\n"
            # 97222 does not have a row of latest data to make sure it still works
            "iso1:us#fips:97222,2020-04-02,No Recent County,county,3\n"
            "iso1:us#fips:97222,2020-04-04,No Recent County,county,5\n"
        )
    )
    assert {region.location_id for region, _ in ts.iter_one_regions()} == {
        "iso1:us#cbsa:10100",
        "iso1:us#fips:97111",
        "iso1:us#fips:97222",
    }
    for it_region, it_one_region in ts.iter_one_regions():
        one_region = ts.get_one_region(it_region)
        assert (one_region.data.fillna("") == it_one_region.data.fillna("")).all(axis=None)
        assert one_region.latest == it_one_region.latest
        assert one_region.provenance == it_one_region.provenance
        assert one_region.region == it_region
        assert one_region.region == it_one_region.region


def test_drop_regions_without_population():
    cbsa_with_pop = Region.from_cbsa_code("10100")
    fips_with_pop = Region.from_fips("97111")
    cbsa_without_pop = Region.from_cbsa_code("20300")
    fips_without_pop = Region.from_fips("97222")
    m1 = FieldName("m1")
    regions_with_pop = [cbsa_with_pop, fips_with_pop]
    all_regions = regions_with_pop + [cbsa_without_pop, fips_without_pop]
    static_populations = {r: {CommonFields.POPULATION: 80_000} for r in regions_with_pop}
    ts_in = test_helpers.build_dataset(
        {r: {m1: [1]} for r in all_regions}, static_by_region_then_field_name=static_populations,
    )
    ts_expected = test_helpers.build_dataset(
        {r: {m1: [1]} for r in regions_with_pop},
        static_by_region_then_field_name=static_populations,
    )
    with structlog.testing.capture_logs() as logs:
        ts_out = timeseries.drop_regions_without_population(
            ts_in, [fips_without_pop.location_id], structlog.get_logger()
        )
    test_helpers.assert_dataset_like(ts_out, ts_expected)

    assert [l["event"] for l in logs] == ["Dropping unexpected regions without populaton"]
    assert [l["location_ids"] for l in logs] == [[cbsa_without_pop.location_id]]


def test_merge_provenance():
    ts = timeseries.MultiRegionDataset.from_csv(
        io.StringIO(
            "location_id,date,county,aggregate_level,m1\n"
            "iso1:us#cbsa:10100,2020-04-02,,,\n"
            "iso1:us#cbsa:10100,2020-04-03,,,\n"
            "iso1:us#cbsa:10100,,,,\n"
            "iso1:us#fips:97111,2020-04-02,Bar County,county,2\n"
            "iso1:us#fips:97111,2020-04-04,Bar County,county,4\n"
            "iso1:us#fips:97111,,Bar County,county,4\n"
        )
    ).add_provenance_csv(
        io.StringIO("location_id,variable,provenance\n" "iso1:us#cbsa:10100,m1,ts110100prov\n")
    )

    with pytest.raises(NotImplementedError):
        ts.add_provenance_csv(
            io.StringIO("location_id,variable,provenance\n" "iso1:us#fips:97111,m1,ts197111prov\n")
        )


def test_append_tags():
    region_sf = Region.from_fips("06075")
    cases_values = [100, 200, 300, 400]
    metrics_sf = {
        CommonFields.POSITIVE_TESTS: TimeseriesLiteral([1, 2, 3, 4], provenance="pt_src2"),
        CommonFields.CASES: cases_values,
    }
    dataset_in = test_helpers.build_dataset({region_sf: metrics_sf})
    tag_sf_cases = test_helpers.make_tag(TagType.CUMULATIVE_TAIL_TRUNCATED, date="2020-04-02")
    tag_df = test_helpers.make_tag_df(
        region_sf, CommonFields.CASES, DemographicBucket.ALL, [tag_sf_cases]
    )
    dataset_out = dataset_in.append_tag_df(tag_df)
    metrics_sf[CommonFields.CASES] = TimeseriesLiteral(cases_values, annotation=[tag_sf_cases])
    dataset_expected = test_helpers.build_dataset({region_sf: metrics_sf})

    test_helpers.assert_dataset_like(dataset_out, dataset_expected)


def test_add_provenance_all_with_tags():
    """Checks that add_provenance_all (and add_provenance_series that it calls) fails when tags
    already exist."""
    region = Region.from_state("TX")
    cases_values = [100, 200, 300, 400]
    timeseries = TimeseriesLiteral(cases_values, annotation=[(test_helpers.make_tag())])
    dataset_in = test_helpers.build_dataset({region: {CommonFields.CASES: timeseries}})

    with pytest.raises(NotImplementedError):
        dataset_in.add_provenance_all("prov_prov")


def test_join_columns_with_tags():
    """Checks that join_columns preserves tags."""
    region = Region.from_state("TX")
    cases_values = [100, 200, 300, 400]
    ts_lit = TimeseriesLiteral(cases_values, annotation=[test_helpers.make_tag()])
    dataset_cases = test_helpers.build_dataset({region: {CommonFields.CASES: ts_lit}})
    dataset_deaths = test_helpers.build_dataset({region: {CommonFields.DEATHS: ts_lit}})

    dataset_out = dataset_cases.join_columns(dataset_deaths)

    assert len(dataset_out.tag) == 2
    # The following checks that the tags in `ts_lit` have been preserved.
    dataset_expected = test_helpers.build_dataset(
        {region: {CommonFields.CASES: ts_lit, CommonFields.DEATHS: ts_lit}}
    )

    test_helpers.assert_dataset_like(dataset_out, dataset_expected)


def test_drop_column_with_tags():
    region = Region.from_state("TX")
    cases_values = [100, 200, 300, 400]
    ts_lit = TimeseriesLiteral(cases_values, annotation=[test_helpers.make_tag()])

    dataset_in = test_helpers.build_dataset(
        {region: {CommonFields.CASES: ts_lit, CommonFields.DEATHS: ts_lit}}
    )

    dataset_out = dataset_in.drop_column_if_present(CommonFields.DEATHS)

    assert len(dataset_out.tag) == 1
    dataset_expected = test_helpers.build_dataset({region: {CommonFields.CASES: ts_lit}})
    test_helpers.assert_dataset_like(dataset_out, dataset_expected)


def test_drop_na_columns():
    tag = test_helpers.make_tag()
    timeseries_real = {
        CommonFields.CASES: TimeseriesLiteral([1, 2], annotation=[tag]),
    }
    static_real = {CommonFields.STAFFED_BEDS: 3}
    ds = test_helpers.build_default_region_dataset(
        # Adds CASES with real values, which won't be dropped, and a tag for DEATHS, that will be
        # dropped.
        {**timeseries_real, CommonFields.DEATHS: TimeseriesLiteral([], annotation=[tag])},
        static=static_real,
    )
    # The test_helper functions don't do a good job of creating fields that are all NA so the
    # following inserts time series DEATHS and static ICU_BEDS, then asserts that they were
    # inserted.
    timeseries_bucketed_with_na = ds.timeseries_bucketed.copy()
    timeseries_bucketed_with_na.loc[:, CommonFields.DEATHS] = np.nan
    static_with_na = ds.static.copy()
    static_with_na.loc[:, CommonFields.ICU_BEDS] = np.nan
    ds = dataclasses.replace(
        ds, timeseries_bucketed=timeseries_bucketed_with_na, static=static_with_na
    )
    assert CommonFields.DEATHS in ds.timeseries_bucketed.columns
    assert CommonFields.ICU_BEDS in ds.static.columns

    dataset_out = ds.drop_na_columns()

    dataset_expected = test_helpers.build_default_region_dataset(
        timeseries_real, static=static_real
    )
    test_helpers.assert_dataset_like(dataset_out, dataset_expected)


def test_drop_na_columns_no_tags():
    timeseries_real = {CommonFields.CASES: [1, 2]}
    tag = test_helpers.make_tag()
    ds = test_helpers.build_default_region_dataset(
        # Add a tag for DEATHS, that will be dropped.
        {**timeseries_real, CommonFields.DEATHS: TimeseriesLiteral([], annotation=[tag])}
    )
    # The test_helper functions don't do a good job of creating fields that are all NA so the
    # following inserts time series DEATHS and static ICU_BEDS, then asserts that they were
    # inserted.
    timeseries_bucketed_with_na = ds.timeseries_bucketed.copy()
    timeseries_bucketed_with_na.loc[:, CommonFields.DEATHS] = np.nan
    ds = dataclasses.replace(ds, timeseries_bucketed=timeseries_bucketed_with_na)
    assert CommonFields.DEATHS in ds.timeseries_bucketed.columns

    dataset_out = ds.drop_na_columns()

    dataset_expected = test_helpers.build_default_region_dataset(timeseries_real)
    test_helpers.assert_dataset_like(dataset_out, dataset_expected)


def test_drop_column_with_tags_and_bucket():
    age_40s = DemographicBucket("age:40-49")
    ts_lit = TimeseriesLiteral([10, 20, 30], annotation=[test_helpers.make_tag()])
    data_cases = {CommonFields.CASES: {age_40s: ts_lit, DemographicBucket.ALL: ts_lit}}
    data_deaths = {CommonFields.DEATHS: {age_40s: ts_lit}}

    dataset_in = test_helpers.build_default_region_dataset({**data_cases, **data_deaths})
    assert len(dataset_in.tag) == 3

    dataset_out = dataset_in.drop_column_if_present(CommonFields.DEATHS)

    assert len(dataset_out.tag) == 2
    dataset_expected = test_helpers.build_default_region_dataset({**data_cases})
    test_helpers.assert_dataset_like(dataset_out, dataset_expected)


def test_timeseries_empty_timeseries_and_static():
    # Check that empty dataset creates a MultiRegionDataset
    # and that get_one_region raises expected exception.
    dataset = timeseries.MultiRegionDataset.new_without_timeseries()
    with pytest.raises(timeseries.RegionLatestNotFound):
        dataset.get_one_region(Region.from_fips("01001"))


def test_timeseries_empty():
    # Check that empty geodata_timeseries_df creates a MultiRegionDataset
    # and that get_one_region raises expected exception.
    dataset = timeseries.MultiRegionDataset.from_timeseries_df(
        pd.DataFrame([], columns=[CommonFields.LOCATION_ID, CommonFields.DATE])
    )
    with pytest.raises(timeseries.RegionLatestNotFound):
        dataset.get_one_region(Region.from_fips("01001"))


def test_timeseries_empty_static_not_empty():
    # Check that empty timeseries does not prevent static data working as expected.
    dataset = timeseries.MultiRegionDataset.from_timeseries_df(
        pd.DataFrame([], columns=[CommonFields.LOCATION_ID, CommonFields.DATE])
    ).add_static_values(pd.DataFrame([{"location_id": "iso1:us#fips:97111", "m1": 1234}]))
    assert dataset.get_one_region(Region.from_fips("97111")).latest["m1"] == 1234


def test_from_timeseries_df_fips_location_id_mismatch():
    df = test_helpers.read_csv_str(
        "                  location_id, fips,      date,m1\n"
        "iso1:us#iso2:us-tx#fips:48197,48201,2020-04-02, 2\n"
        "iso1:us#iso2:us-tx#fips:48201,48201,2020-04-02, 2\n",
        skip_spaces=True,
    )
    with pytest.warns(timeseries.ExtraColumnWarning, match="48201"):
        timeseries.MultiRegionDataset.from_timeseries_df(df)


def test_from_timeseries_df_no_fips_no_warning():
    df = test_helpers.read_csv_str(
        "            location_id, fips,      date,m1\n"
        "                iso1:us,     ,2020-04-02, 2\n",
        skip_spaces=True,
    )
    timeseries.MultiRegionDataset.from_timeseries_df(df)


def test_from_timeseries_df_fips_state_mismatch():
    df = test_helpers.read_csv_str(
        "                  location_id,state,      date,m1\n"
        "iso1:us#iso2:us-tx#fips:48197,   TX,2020-04-02, 2\n"
        "iso1:us#iso2:us-tx#fips:48201,   IL,2020-04-02, 2\n",
        skip_spaces=True,
    )
    with pytest.warns(timeseries.ExtraColumnWarning, match="48201"):
        timeseries.MultiRegionDataset.from_timeseries_df(df)


def test_from_timeseries_df_bad_level():
    df = test_helpers.read_csv_str(
        "                  location_id, aggregate_level,      date,m1\n"
        "iso1:us#iso2:us-tx#fips:48201,          county,2020-04-02, 2\n"
        "iso1:us#iso2:us-tx#fips:48197,           state,2020-04-02, 2\n"
        "           iso1:us#iso2:us-tx,           state,2020-04-02, 2\n",
        skip_spaces=True,
    )
    with pytest.warns(timeseries.ExtraColumnWarning, match="48197"):
        timeseries.MultiRegionDataset.from_timeseries_df(df)


def test_combined_timeseries():
    ds1 = timeseries.MultiRegionDataset.from_csv(
        io.StringIO(
            "location_id,date,county,aggregate_level,m1\n"
            "iso1:us#cbsa:10100,2020-04-02,,,2.2\n"
            "iso1:us#cbsa:10100,2020-04-03,,,3.3\n"
            "iso1:us#fips:97111,2020-04-02,Bar County,county,2\n"
            "iso1:us#fips:97111,2020-04-04,Bar County,county,4\n"
        )
    ).add_provenance_csv(
        io.StringIO("location_id,variable,provenance\n" "iso1:us#cbsa:10100,m1,ds110100prov\n")
    )
    ds2 = timeseries.MultiRegionDataset.from_csv(
        io.StringIO(
            "location_id,date,county,aggregate_level,m1\n"
            "iso1:us#cbsa:10100,2020-04-02,,,333\n"
            "iso1:us#cbsa:10100,2020-04-03,,,333\n"
            "iso1:us#fips:97222,2020-04-03,Foo County,county,30\n"
            "iso1:us#fips:97222,2020-04-04,Foo County,county,40\n"
        )
    ).add_provenance_csv(
        io.StringIO("location_id,variable,provenance\n" "iso1:us#cbsa:10100,m1,ds110100prov\n")
    )
    combined = timeseries.combined_datasets({FieldName("m1"): [ds1, ds2]}, {})
    expected = timeseries.MultiRegionDataset.from_csv(
        io.StringIO(
            "location_id,date,m1\n"
            "iso1:us#cbsa:10100,2020-04-02,2.2\n"
            "iso1:us#cbsa:10100,2020-04-03,3.3\n"
            "iso1:us#fips:97111,2020-04-02,2\n"
            "iso1:us#fips:97111,2020-04-04,4\n"
            "iso1:us#fips:97222,2020-04-03,30\n"
            "iso1:us#fips:97222,2020-04-04,40\n"
        )
    ).add_provenance_csv(
        io.StringIO("location_id,variable,provenance\n" "iso1:us#cbsa:10100,m1,ds110100prov\n")
    )

    test_helpers.assert_dataset_like(expected, combined)


def test_combined_annotation():
    ts1a = TimeseriesLiteral(
        [0, 2, 4],
        annotation=[
            test_helpers.make_tag(date="2020-04-01"),
            test_helpers.make_tag(date="2020-04-02"),
        ],
    )
    ts1b = [100, 200, 300]
    ds1 = test_helpers.build_default_region_dataset(
        {CommonFields.ICU_BEDS: ts1a, CommonFields.CASES: ts1b}
    )
    ts2a = TimeseriesLiteral([1, 3, 5], annotation=[test_helpers.make_tag(date="2020-04-01")],)
    ts2b = [150, 250, 350]
    ds2 = test_helpers.build_default_region_dataset(
        {CommonFields.ICU_BEDS: ts2a, CommonFields.CASES: ts2b}
    )
    combined = timeseries.combined_datasets(
        {CommonFields.ICU_BEDS: [ds1, ds2], CommonFields.CASES: [ds2, ds1]}, {}
    )

    expected = test_helpers.build_default_region_dataset(
        {CommonFields.ICU_BEDS: ts1a, CommonFields.CASES: ts2b}
    )

    test_helpers.assert_dataset_like(combined, expected)


def test_combined_missing_field():
    ts1 = timeseries.MultiRegionDataset.from_csv(
        io.StringIO(
            "location_id,date,county,aggregate_level,m1\n"
            "iso1:us#fips:97111,2020-04-02,Bar County,county,2\n"
            "iso1:us#fips:97111,2020-04-04,Bar County,county,4\n"
        )
    )
    ts2 = timeseries.MultiRegionDataset.from_csv(
        io.StringIO(
            "location_id,date,county,aggregate_level,m2\n"
            "iso1:us#fips:97111,2020-04-02,Bar County,county,111\n"
            "iso1:us#fips:97111,2020-04-04,Bar County,county,111\n"
        )
    )
    # m1 is output, m2 is dropped.
    field_source_map = {FieldName("m1"): [ts1, ts2]}

    # Check that combining finishes and produces the expected result.
    combined_1 = timeseries.combined_datasets(field_source_map, {})
    expected = timeseries.MultiRegionDataset.from_csv(
        io.StringIO(
            "location_id,date,m1\n"
            "iso1:us#fips:97111,2020-04-02,2\n"
            "iso1:us#fips:97111,2020-04-04,4\n"
        )
    )
    test_helpers.assert_dataset_like(expected, combined_1)

    # Because there is only one source for the output timeseries reversing the source list
    # produces the same output.
    combined_2 = timeseries.combined_datasets(
        {name: list(reversed(source_list)) for name, source_list in field_source_map.items()}, {}
    )
    test_helpers.assert_dataset_like(expected, combined_2)


def test_combined_static():
    ds1 = timeseries.MultiRegionDataset.from_csv(
        io.StringIO(
            "location_id,date,county,aggregate_level,s1\n"
            "iso1:us#cbsa:10100,,,,\n"
            "iso1:us#fips:97222,,Foo County,county,22\n"
        )
    )
    ds2 = timeseries.MultiRegionDataset.from_csv(
        io.StringIO(
            "location_id,date,county,aggregate_level,s1\n"
            "iso1:us#cbsa:10100,,,,111\n"
            "iso1:us#fips:97222,,Foo County,county,222\n"
        )
    )
    combined = timeseries.combined_datasets({}, {FieldName("s1"): [ds1, ds2]})
    expected = timeseries.MultiRegionDataset.from_csv(
        io.StringIO("location_id,date,s1\n" "iso1:us#cbsa:10100,,111\n" "iso1:us#fips:97222,,22\n")
    )

    test_helpers.assert_dataset_like(expected, combined, drop_na_timeseries=True)


def test_timeseries_rows():
    ts = timeseries.MultiRegionDataset.from_csv(
        io.StringIO(
            "location_id,county,aggregate_level,date,m1,m2,population\n"
            "iso1:us#iso2:us-tx,Texas,state,2020-04-01,4,2,\n"
            "iso1:us#iso2:us-tx,Texas,state,2020-04-02,4,4,\n"
            "iso1:us#iso2:us-tx,Texas,state,,,,2500\n"
            "iso1:us#iso2:us-az,Arizona,state,2020-04-01,8,20,\n"
            "iso1:us#iso2:us-az,Arizona,state,2020-04-02,12,40,\n"
            "iso1:us#iso2:us-az,Arizona,state,,,,7500\n"
        )
    )

    rows = ts.timeseries_rows()
    expected = test_helpers.read_csv_str(
        "       location_id,variable,demographic_bucket,2020-04-02,2020-04-01\n"
        "iso1:us#iso2:us-az,      m1,               all,        12,         8\n"
        "iso1:us#iso2:us-az,      m2,               all,        40,        20\n"
        "iso1:us#iso2:us-tx,      m1,               all,         4,         4\n"
        "iso1:us#iso2:us-tx,      m2,               all,         4,         2\n",
        skip_spaces=True,
    ).set_index([CommonFields.LOCATION_ID, PdFields.VARIABLE, PdFields.DEMOGRAPHIC_BUCKET])
    pd.testing.assert_frame_equal(rows, expected, check_dtype=False, check_exact=False)


def test_multi_region_dataset_get_subset():
    region_us = Region.from_iso1("us")
    region_tx = Region.from_state("TX")
    region_county = Region.from_fips("97222")
    region_cbsa = Region.from_cbsa_code("10100")
    m1 = FieldName("m1")
    m2 = FieldName("m2")
    ds = test_helpers.build_dataset(
        {
            region_us: {m1: [100], m2: [200]},
            region_tx: {m1: [4], m2: [2]},
            region_county: {m1: [1], m2: [2]},
            region_cbsa: {m1: [1], m2: [2], CommonFields.POPULATION: [20_000]},
        },
        static_by_region_then_field_name={
            region_us: {CommonFields.POPULATION: 10_000},
            region_tx: {CommonFields.POPULATION: 5_000},
            region_county: {CommonFields.POPULATION: 1_000},
        },
    )

    subset = ds.get_subset(aggregation_level=AggregationLevel.COUNTRY)
    assert subset.static.at["iso1:us", CommonFields.POPULATION] == 10000

    subset = ds.get_subset(fips="97222")
    assert subset.timeseries.at[("iso1:us#fips:97222", "2020-04-01"), "m2"] == 2

    subset = ds.get_subset(state="TX")
    assert subset.static.at["iso1:us#iso2:us-tx", CommonFields.POPULATION] == 5000

    subset = ds.get_subset(states=["TX"])
    assert subset.static.at["iso1:us#iso2:us-tx", CommonFields.POPULATION] == 5000

    subset = ds.get_subset(location_id_matches=r"\A(iso1\:us|iso1\:us\#cbsa.+)\Z")
    assert {r.location_id for r, _ in subset.iter_one_regions()} == {
        "iso1:us",
        "iso1:us#cbsa:10100",
    }


def test_multi_region_dataset_get_subset_with_buckets():
    # Make some regions at different levels
    region_us = Region.from_iso1("us")
    region_tx = Region.from_state("TX")
    region_la = Region.from_fips("06037")
    age_40s = DemographicBucket("age:40-49")
    data_us = {region_us: {CommonFields.CASES: [100, 200]}}
    data_tx = {region_tx: {CommonFields.CASES: [10, 20]}}
    data_la = {region_la: {CommonFields.CASES: {DemographicBucket.ALL: [5, 10], age_40s: [1, 2]}}}
    ds = test_helpers.build_dataset({**data_us, **data_tx, **data_la})

    ds_expected = test_helpers.build_dataset({**data_us, **data_la})
    test_helpers.assert_dataset_like(ds.get_regions_subset([region_us, region_la]), ds_expected)
    test_helpers.assert_dataset_like(ds.partition_by_region(exclude=[region_tx])[0], ds_expected)


def test_write_read_dataset_pointer_with_source_url(tmpdir):
    pointer = _make_dataset_pointer(tmpdir)
    url_str1 = UrlStr("http://foo.com/1")
    url_str2 = UrlStr("http://foo.com/2")
    url_str3 = UrlStr("http://foo.com/3")

    ts1a = TimeseriesLiteral(
        [0, 2, 4],
        annotation=[
            test_helpers.make_tag(date="2020-04-01"),
            test_helpers.make_tag(date="2020-04-02"),
        ],
        source_url=url_str1,
    )
    ts1b = TimeseriesLiteral([100, 200, 300], source_url=[url_str2, url_str3])
    dataset_in = test_helpers.build_default_region_dataset(
        {CommonFields.ICU_BEDS: ts1a, CommonFields.CASES: ts1b}
    )

    dataset_in.write_to_dataset_pointer(pointer)

    dataset_read = timeseries.MultiRegionDataset.read_from_pointer(pointer)

    test_helpers.assert_dataset_like(dataset_read, dataset_in)
    source_url_read = dataset_read.get_one_region(test_helpers.DEFAULT_REGION).source_url
    assert source_url_read[CommonFields.ICU_BEDS] == [url_str1]
    # Copy to a set because the order of the URLs in the source_url may change.
    assert set(source_url_read[CommonFields.CASES]) == {url_str2, url_str3}


def test_pickle():
    ts = TimeseriesLiteral(
        [0, 2, 4],
        annotation=[
            test_helpers.make_tag(date="2020-04-01"),
            test_helpers.make_tag(date="2020-04-02"),
        ],
        source_url=UrlStr("http://public.com"),
    )
    ds_in = test_helpers.build_default_region_dataset({CommonFields.CASES: ts})

    ds_out = pickle.loads(pickle.dumps(ds_in))

    test_helpers.assert_dataset_like(ds_in, ds_out)


def test_make_source_tags():
    url_str = UrlStr("http://foo.com/1")

    ts_prov_only = TimeseriesLiteral(
        [0, 2, 4], annotation=[test_helpers.make_tag(date="2020-04-01"),], provenance="prov_only",
    )
    ts_with_url = TimeseriesLiteral([3, 5, 7], provenance="prov_with_url", source_url=url_str)
    dataset_in = test_helpers.build_default_region_dataset(
        {CommonFields.ICU_BEDS: ts_prov_only, CommonFields.CASES: ts_with_url}
    )

    dataset_out = timeseries.make_source_tags(dataset_in)

    source_tag_prov_only = taglib.Source("prov_only")
    ts_prov_only_expected = TimeseriesLiteral(
        [0, 2, 4],
        annotation=[test_helpers.make_tag(date="2020-04-01"),],
        source=source_tag_prov_only,
    )
    source_tag_prov_with_url = taglib.Source("prov_with_url", url=url_str)
    ts_with_url_expected = TimeseriesLiteral([3, 5, 7], source=source_tag_prov_with_url,)
    dataset_expected = test_helpers.build_default_region_dataset(
        {CommonFields.ICU_BEDS: ts_prov_only_expected, CommonFields.CASES: ts_with_url_expected}
    )
    test_helpers.assert_dataset_like(dataset_out, dataset_expected)

    one_region = dataset_out.get_one_region(test_helpers.DEFAULT_REGION)
    assert one_region.sources_all_bucket(CommonFields.ICU_BEDS) == [source_tag_prov_only]
    assert one_region.sources_all_bucket(CommonFields.CASES) == [source_tag_prov_with_url]


def test_make_source_tags_no_urls():
    # There was a bug where `./run.py data update` failed at the very end when no timeseries had
    # a source_url. This tests for it.
    ts_prov_only = TimeseriesLiteral(
        [0, 2, 4], annotation=[test_helpers.make_tag(date="2020-04-01"),], provenance="prov_only",
    )
    dataset_in = test_helpers.build_default_region_dataset({CommonFields.ICU_BEDS: ts_prov_only})

    dataset_out = timeseries.make_source_tags(dataset_in)

    source_tag_prov_only = taglib.Source("prov_only")
    ts_prov_only_expected = TimeseriesLiteral(
        [0, 2, 4],
        annotation=[test_helpers.make_tag(date="2020-04-01"),],
        source=source_tag_prov_only,
    )
    dataset_expected = test_helpers.build_default_region_dataset(
        {CommonFields.ICU_BEDS: ts_prov_only_expected}
    )
    test_helpers.assert_dataset_like(dataset_out, dataset_expected)

    one_region = dataset_out.get_one_region(test_helpers.DEFAULT_REGION)
    assert one_region.sources_all_bucket(CommonFields.ICU_BEDS) == [source_tag_prov_only]


def test_make_source_url_tags():
    url_str = UrlStr("http://foo.com/1")

    source_tag_prov_only = taglib.Source("prov_only")
    ts_prov_only = TimeseriesLiteral(
        [0, 2, 4],
        annotation=[test_helpers.make_tag(date="2020-04-01"),],
        source=source_tag_prov_only,
    )
    source_tag_prov_with_url = taglib.Source("prov_with_url", url=url_str)
    ts_with_url = TimeseriesLiteral([3, 5, 7], source=source_tag_prov_with_url,)
    dataset_in = test_helpers.build_default_region_dataset(
        {CommonFields.ICU_BEDS: ts_prov_only, CommonFields.CASES: ts_with_url}
    )

    dataset_out = timeseries.make_source_url_tags(dataset_in)

    ts_with_url_expected = TimeseriesLiteral(
        [3, 5, 7], source=source_tag_prov_with_url, source_url=url_str
    )
    dataset_expected = test_helpers.build_default_region_dataset(
        {CommonFields.ICU_BEDS: ts_prov_only, CommonFields.CASES: ts_with_url_expected}
    )
    test_helpers.assert_dataset_like(dataset_out, dataset_expected)


def test_make_source_url_tags_no_source_tags():
    dataset_in = test_helpers.build_default_region_dataset({CommonFields.CASES: [1, 2, 3]})
    dataset_out = timeseries.make_source_url_tags(dataset_in)
    assert dataset_in == dataset_out


def test_make_source_url_tags_has_source_url():
    url_str = UrlStr("http://foo.com/1")
    dataset_in = test_helpers.build_default_region_dataset(
        {CommonFields.CASES: TimeseriesLiteral([1, 2, 3], source_url=url_str)}
    )
    with pytest.raises(AssertionError):
        timeseries.make_source_url_tags(dataset_in)


def test_check_timeseries_structure_empty():
    timeseries._check_timeseries_wide_vars_structure(
        timeseries.EMPTY_TIMESERIES_WIDE_VARIABLES_DF, bucketed=False
    )
    timeseries._check_timeseries_wide_vars_structure(
        timeseries.EMPTY_TIMESERIES_BUCKETED_WIDE_VARIABLES_DF, bucketed=True
    )


def test_make_and_pickle_demographic_data():
    location_id = test_helpers.DEFAULT_REGION.location_id
    date_0 = test_helpers.DEFAULT_START_DATE
    date_1 = pd.to_datetime(test_helpers.DEFAULT_START_DATE) + pd.to_timedelta(1, unit="day")
    m1 = FieldName("m1")
    age20s = DemographicBucket("age:20-29")
    age30s = DemographicBucket("age:30-39")
    all = DemographicBucket("all")

    ds = test_helpers.build_default_region_dataset(
        {m1: {age20s: [1, 2, 3], age30s: [5, 6, 7], all: [8, 9, None]}}
    )

    assert ds.timeseries_bucketed.at[(location_id, age30s, date_0), m1] == 5
    assert ds.timeseries_bucketed.at[(location_id, all, date_1), m1] == 9

    ds_unpickled = pickle.loads(pickle.dumps(ds))

    test_helpers.assert_dataset_like(ds, ds_unpickled)


def test_combine_demographic_data_basic():
    m1 = FieldName("m1")
    age20s = DemographicBucket("age:20-29")
    age30s = DemographicBucket("age:30-39")
    age40s = DemographicBucket("age:40-49")
    ds1 = test_helpers.build_default_region_dataset(
        {m1: {age20s: [21, 22, 23], age30s: [31, 32, 33],}}
    )
    ds2 = test_helpers.build_default_region_dataset(
        {m1: {age30s: [32, 33, 34], age40s: [42, 43, 44],}}
    )

    combined = timeseries.combined_datasets({m1: [ds1, ds2]}, {})
    test_helpers.assert_dataset_like(combined, ds1)

    combined = timeseries.combined_datasets({m1: [ds2, ds1]}, {})
    test_helpers.assert_dataset_like(combined, ds2)


def test_combine_demographic_data_multiple_distributions():
    """All time-series within a variable are treated as a unit when combining"""
    m1 = FieldName("m1")
    m2 = FieldName("m2")
    all = DemographicBucket("all")
    age_20s = DemographicBucket("age:20-29")
    age_30s = DemographicBucket("age:30-39")
    region_ak = Region.from_state("AK")
    region_ca = Region.from_state("CA")

    ds1 = test_helpers.build_dataset(
        {
            region_ak: {m1: {all: TimeseriesLiteral([1.1, 2.1], provenance="ds1_ak_m1_all")}},
            region_ca: {m1: {age_20s: TimeseriesLiteral([2.1, 3.1], provenance="ds1_ca_m1_20s")}},
        }
    )

    ds2 = test_helpers.build_dataset(
        {
            region_ak: {m1: {all: TimeseriesLiteral([1, 2], provenance="ds2_ak_m1_all")}},
            region_ca: {
                m1: {
                    age_30s: TimeseriesLiteral([3, 4], provenance="ds2_ca_m1_30s"),
                    all: TimeseriesLiteral([5, 6], provenance="ds2_ca_m1_all"),
                },
                m2: {age_30s: TimeseriesLiteral([6, 7], provenance="ds2_ca_m2_30s")},
            },
        }
    )

    combined = timeseries.combined_datasets({m1: [ds1, ds2], m2: [ds1, ds2]}, {})

    ds_expected = test_helpers.build_dataset(
        {
            region_ak: {m1: {all: TimeseriesLiteral([1.1, 2.1], provenance="ds1_ak_m1_all")}},
            region_ca: {
                m1: {
                    age_20s: TimeseriesLiteral([2.1, 3.1], provenance="ds1_ca_m1_20s"),
                    all: TimeseriesLiteral([5, 6], provenance="ds2_ca_m1_all"),
                },
                m2: {age_30s: TimeseriesLiteral([6, 7], provenance="ds2_ca_m2_30s")},
            },
        }
    )
    test_helpers.assert_dataset_like(combined, ds_expected)


def test_bucketed_latest_missing_location_id(nyc_region: Region):
    dataset = test_helpers.build_default_region_dataset({CommonFields.CASES: [1, 2, 3]})
    # nyc_region = Region.from_fips("97222")
    output = dataset._bucketed_latest_for_location_id(nyc_region.location_id)
    expected = pd.DataFrame(
        [],
        index=pd.MultiIndex.from_tuples([], names=[PdFields.DEMOGRAPHIC_BUCKET]),
        columns=pd.Index([CommonFields.CASES], name="variable"),
        dtype="float",
    )
    pd.testing.assert_frame_equal(expected, output)


def test_bucketed_latest(nyc_region: Region):
    m1 = FieldName("m1")
    age20s = DemographicBucket("age:20-29")
    age30s = DemographicBucket("age:30-39")

    dataset = test_helpers.build_default_region_dataset(
        {m1: {age20s: [21, 22, 23], age30s: [31, 32, 33],}}
    )
    bucketed_latest = dataset._bucketed_latest_for_location_id(
        test_helpers.DEFAULT_REGION.location_id
    )
    expected = pd.DataFrame(
        [{"m1": 23}, {"m1": 33}],
        index=pd.Index([age20s, age30s], name=PdFields.DEMOGRAPHIC_BUCKET),
        columns=pd.Index([m1], name="variable"),
    )
    pd.testing.assert_frame_equal(bucketed_latest, expected)


def test_one_region_demographic_distributions():
    m1 = FieldName("m1")
    age20s = DemographicBucket("age:20-29")
    age30s = DemographicBucket("age:30-39")
    dataset = test_helpers.build_default_region_dataset(
        {m1: {age20s: [21, 22, 23], age30s: [31, 32, 33], DemographicBucket.ALL: [20, 21, 22]}}
    )
    one_region = dataset.get_one_region(test_helpers.DEFAULT_REGION)

    expected = {m1: {"age": {"20-29": 23, "30-39": 33}}}
    assert one_region.demographic_distributions_by_field == expected


def test_one_region_demographic_distributions_overlapping_buckets():
    m1 = FieldName("m1")
    m2 = FieldName("m2")
    age20s = DemographicBucket("age:20-29")
    age30s = DemographicBucket("age:30-39")
    # Presumably 25 to 29 is from a different age distribution as it overlaps with age bucket above.
    # Make sure that different age bucketing doesn't polute other variables.
    age25to29 = DemographicBucket("age:25-29")

    dataset = test_helpers.build_default_region_dataset(
        {
            m1: {age20s: [21, 22, 23], age30s: [31, 32, 33], DemographicBucket.ALL: [20, 21, 22]},
            m2: {DemographicBucket.ALL: [20, 21, 22], age25to29: [20, 21, 22]},
        },
    )
    one_region = dataset.get_one_region(test_helpers.DEFAULT_REGION)
    expected = {m1: {"age": {"20-29": 23, "30-39": 33}}, m2: {"age": {"25-29": 22}}}

    assert one_region.demographic_distributions_by_field == expected


def test_print_stats():
    all_bucket = DemographicBucket("all")
    age_20s = DemographicBucket("age:20-29")
    age_30s = DemographicBucket("age:30-39")

    test_helpers.build_default_region_dataset(
        {
            CommonFields.ICU_BEDS: TimeseriesLiteral(
                [0, 2, 4], annotation=[test_helpers.make_tag(date="2020-04-01"),],
            ),
            CommonFields.CASES: [100, 200, 300],
        }
    ).print_stats("DS1")

    test_helpers.build_default_region_dataset(
        {
            CommonFields.CASES: {
                age_20s: TimeseriesLiteral([3, 4, 5], source=taglib.Source(type="MySource")),
                age_30s: [4, 5, 6],
                all_bucket: [1, 2, 3],
            }
        }
    ).print_stats("DS2")


def test_static_and_geo_data():
    region_chi = Region.from_fips("17031")
    ds = test_helpers.build_default_region_dataset(
        {CommonFields.CASES: [0]}, static={CommonFields.POPULATION: 5}, region=region_chi
    )
    assert ds.static_and_geo_data.loc[region_chi.location_id, CommonFields.COUNTY] == "Cook County"
    assert ds.static_and_geo_data.loc[region_chi.location_id, CommonFields.POPULATION] == 5


def test_add_tag_all_bucket():
    region_tx = Region.from_state("TX")
    region_la = Region.from_fips("06037")
    age_40s = DemographicBucket("age:40-49")
    data_tx = {region_tx: {CommonFields.CASES: [10, 20]}}
    data_la = {region_la: {CommonFields.CASES: {DemographicBucket.ALL: [5, 10], age_40s: [1, 2]}}}

    tag = test_helpers.make_tag(date="2020-04-01")
    ds = test_helpers.build_dataset({**data_tx, **data_la}).add_tag_all_bucket(tag)

    expected_tx = {region_tx: {CommonFields.CASES: TimeseriesLiteral([10, 20], annotation=[tag])}}
    expected_la = {
        region_la: {
            CommonFields.CASES: {
                DemographicBucket.ALL: TimeseriesLiteral([5, 10], annotation=[tag]),
                age_40s: [1, 2],
            }
        }
    }
    ds_expected = test_helpers.build_dataset({**expected_tx, **expected_la})
    test_helpers.assert_dataset_like(ds, ds_expected)


def test_add_tag_without_timeseries(tmpdir):
    """Create a dataset with a tag for a timeseries that doesn't exist."""
    pointer = _make_dataset_pointer(tmpdir)

    region_tx = Region.from_state("TX")
    region_la = Region.from_fips("06037")
    data_tx = {region_tx: {CommonFields.CASES: [10, 20]}}

    tag_collection = taglib.TagCollection()
    tag = test_helpers.make_tag(date="2020-04-01")
    tag_collection.add(
        tag,
        location_id=region_la.location_id,
        variable=CommonFields.CASES,
        bucket=DemographicBucket.ALL,
    )

    dataset = test_helpers.build_dataset({**data_tx}).append_tag_df(tag_collection.as_dataframe())

    # Check that the tag was created for region_la, which doesn't have any timeseries data.
    assert set(
        dataset.tag_objects_series.xs(region_la.location_id, level=CommonFields.LOCATION_ID)
    ) == {tag}

    # Check that tag location_id are included in location_ids property.
    assert set(dataset.location_ids) == {region_la.location_id, region_tx.location_id}

    # Check that the tag still exists after writing and reading from disk.
    dataset.write_to_dataset_pointer(pointer)
    dataset_read = timeseries.MultiRegionDataset.read_from_pointer(pointer)
    test_helpers.assert_dataset_like(dataset, dataset_read)


def test_variables():
    # Make a dataset with CASES, DEATHS and ICU_BEDS each appearing in only one of timeseries,
    # static and tag data. This make sure variable names are merged from all three places.
    region_97111 = Region.from_fips("97111")
    tag_collection = taglib.TagCollection()
    tag_collection.add(
        test_helpers.make_tag(),
        location_id=region_97111.location_id,
        variable=CommonFields.DEATHS,
        bucket=DemographicBucket.ALL,
    )
    ds = test_helpers.build_dataset(
        {region_97111: {CommonFields.CASES: [1, 2, None]}},
        static_by_region_then_field_name={region_97111: {CommonFields.ICU_BEDS: 10}},
    ).append_tag_df(tag_collection.as_dataframe())
    assert set(ds.variables.to_list()) == {
        CommonFields.CASES,
        CommonFields.ICU_BEDS,
        CommonFields.DEATHS,
    }


def test_variables_empty():
    assert timeseries.MultiRegionDataset.new_without_timeseries().variables.to_list() == []


def test_static_long():
    region_cbsa = Region.from_cbsa_code("10100")
    region_fips = Region.from_fips("97111")
    m1 = FieldName("m1")
    ds = test_helpers.build_dataset(
        {},
        static_by_region_then_field_name={
            region_fips: {CommonFields.CAN_LOCATION_PAGE_URL: "http://can.do", m1: 4},
            region_cbsa: {CommonFields.CASES: 3},
        },
    )
    # Use loc[level0].at[level1] as work-around for
    # https://github.com/pandas-dev/pandas/issues/26989
    # TODO(tom): Change to `at[level0, level1]` after upgrading to Pandas >=1.1
    assert (
        ds.static_long.loc[region_fips.location_id].at[CommonFields.CAN_LOCATION_PAGE_URL]
        == "http://can.do"
    )
    assert ds.static_long.loc[region_fips.location_id].at[m1] == 4
    assert ds.static_long.loc[region_cbsa.location_id].at[CommonFields.CASES] == 3

    ds_empty_static = timeseries.MultiRegionDataset.new_without_timeseries()
    assert ds_empty_static.static_long.empty
    assert ds_empty_static.static_long.name == ds.static_long.name
    assert ds_empty_static.static_long.index.names == ds.static_long.index.names


def test_delta_timeseries_removed():
    # This tests time series being removed only, not tags or static values.
    region_tx = Region.from_state("TX")
    region_la = Region.from_fips("06037")
    age_40s = DemographicBucket("age:40-49")
    data_tx = {region_tx: {CommonFields.CASES: [10, 20]}}
    data_la_a = {region_la: {CommonFields.CASES: {DemographicBucket.ALL: [5, 10], age_40s: [1, 2]}}}

    ds_a = test_helpers.build_dataset({**data_tx, **data_la_a})

    data_la_b = {region_la: {CommonFields.CASES: {DemographicBucket.ALL: [5, 10]}}}
    ds_b = test_helpers.build_dataset({**data_tx, **data_la_b})

    delta = timeseries.MultiRegionDatasetDiff(old=ds_a, new=ds_b)
    ds_out = delta.timeseries_removed

    ds_expected = test_helpers.build_dataset({region_la: {CommonFields.CASES: {age_40s: [1, 2]}}})

    test_helpers.assert_dataset_like(ds_out, ds_expected)


def test_drop_observations_after():
    age_40s = DemographicBucket("age:40-49")
    ds_in = test_helpers.build_default_region_dataset(
        {
            CommonFields.CASES: {DemographicBucket.ALL: [5, 10], age_40s: [1, 2, 3]},
            # Check that observation is dropped even when not a True value (ie 0).
            CommonFields.DEATHS: [0, 0, 0],
            # Check what happens when there are no real valued observations after dropping,
            # though the behaviour probably doesn't matter.
            CommonFields.ICU_BEDS: [None, None, 10],
        }
    )

    ds_out = timeseries.drop_observations(ds_in, after=datetime.date(2020, 4, 2))

    tag = test_helpers.make_tag(taglib.TagType.DROP_FUTURE_OBSERVATION, after="2020-04-02")
    ds_expected = test_helpers.build_default_region_dataset(
        {
            CommonFields.CASES: {
                DemographicBucket.ALL: [5, 10],
                age_40s: TimeseriesLiteral([1, 2], annotation=[tag]),
            },
            CommonFields.DEATHS: TimeseriesLiteral([0, 0], annotation=[tag]),
            CommonFields.ICU_BEDS: TimeseriesLiteral([], annotation=[tag]),
        }
    )

    test_helpers.assert_dataset_like(ds_out, ds_expected)


def test_pickle_test_dataset_size(tmp_path: pathlib.Path):
    pkl_path = tmp_path / "testfile.pkl.gz"
    test_dataset = test_helpers.load_test_dataset()
    test_dataset.get_timeseries_not_bucketed_wide_dates(CommonFields.CASES)
    test_dataset.to_compressed_pickle(pkl_path)
    assert pkl_path.stat().st_size < 800_000

    loaded_dataset = timeseries.MultiRegionDataset.from_compressed_pickle(pkl_path)

    test_helpers.assert_dataset_like(test_dataset, loaded_dataset)

pandas.to_datetime pandas.testing.assert_frame_equal pandas.Index pandas.DataFrame pandas.to_timedelta pandas.MultiIndex.from_tuples pandas.testing.assert_series_equal pandas.Series