# Licensed under a 3-clause BSD style license - see LICENSE.rst

"""
This module tests some of the methods related to the ``ECSV``
reader/writer.
"""

import copy
import io
import os
import sys
from contextlib import nullcontext
from io import StringIO
from pprint import pformat

import numpy as np
import pytest
import yaml
from numpy.testing import assert_array_equal

from astropy import units as u
from astropy.io import ascii
from astropy.io.ascii.ecsv import DELIMITERS, InvalidEcsvDatatypeWarning
from astropy.io.ascii.tests.common import TEST_DIR
from astropy.io.misc.ecsv import ECSVEngine
from astropy.io.tests.mixin_columns import compare_attrs, mixin_cols, serialized_names
from astropy.table import Column, QTable, Table
from astropy.table.column import MaskedColumn
from astropy.table.table_helpers import simple_table
from astropy.time import Time
from astropy.units import QuantityInfo
from astropy.units import allclose as quantity_allclose
from astropy.utils.compat.optional_deps import (
    HAS_BZ2,
    HAS_LZMA,
    HAS_PANDAS,
    HAS_PYARROW,
)
from astropy.utils.masked import Masked

ENGINE_PARAMS = [
    {"format": "ascii.ecsv"},
    {"format": "ecsv", "engine": "io.ascii"},
]
if HAS_PANDAS:
    ENGINE_PARAMS.append({"format": "ecsv", "engine": "pandas"})
if HAS_PYARROW:
    ENGINE_PARAMS.append({"format": "ecsv", "engine": "pyarrow"})


@pytest.fixture(scope="module", params=ENGINE_PARAMS)
def format_engine(request):
    return request.param


def patch_format_write(format_engine):
    """Return a compatible format_engine dict for available writers"""
    out = format_engine.copy()
    if out["format"] == "ecsv":
        # "io.ascii" is the only supported engine for writing via "ecsv" format.
        out["engine"] = "io.ascii"
    return out


DTYPES = [
    "bool",
    "int8",
    "int16",
    "int32",
    "int64",
    "uint8",
    "uint16",
    "uint32",
    "uint64",
    "float16",
    "float32",
    "float64",
    "float128",
    "str",
]
if not hasattr(np, "float128") or os.name == "nt" or sys.maxsize <= 2**32:
    DTYPES.remove("float128")

T_DTYPES = Table()

for dtype in DTYPES:
    if dtype == "bool":
        data = np.array([False, True, False])
    elif dtype == "str":
        data = np.array(["ab 0", "ab, 1", "ab2"])
    else:
        data = np.arange(3, dtype=dtype)
    c = Column(data, unit="m / s", description="descr_" + dtype)
    # Add meta in way that uses the default_factory type and not in alphabetical order
    c.meta["meta " + dtype] = 1
    c.meta["a"] = 2
    T_DTYPES[dtype] = c

# Add meta in way that uses the default_factory type and not in alphabetical order
T_DTYPES.meta["comments"] = ["comment1", "comment2"]
T_DTYPES.meta["a"] = 3

# Corresponds to simple_table()
SIMPLE_LINES = [
    "# %ECSV 1.0",
    "# ---",
    "# datatype:",
    "# - {name: a, datatype: int64}",
    "# - {name: b, datatype: float64}",
    "# - {name: c, datatype: string}",
    "# schema: astropy-2.0",
    "a b c",
    "1 1.0 c",
    "2 2.0 d",
    "3 3.0 e",
]


def test_write_simple():
    """
    Write a simple table with common types.  This shows the compact version
    of serialization with one line per column.
    """
    t = simple_table()

    out = StringIO()
    t.write(out, format="ascii.ecsv")
    assert out.getvalue().splitlines() == SIMPLE_LINES


def test_write_full():
    """
    Write a full-featured table with common types and explicitly checkout output
    """
    t = T_DTYPES["bool", "int64", "float64", "str"]
    lines = [
        "# %ECSV 1.0",
        "# ---",
        "# datatype:",
        "# - name: bool",
        "#   unit: m / s",
        "#   datatype: bool",
        "#   description: descr_bool",
        "#   meta: !!omap",
        "#   - {meta bool: 1}",
        "#   - {a: 2}",
        "# - name: int64",
        "#   unit: m / s",
        "#   datatype: int64",
        "#   description: descr_int64",
        "#   meta: !!omap",
        "#   - {meta int64: 1}",
        "#   - {a: 2}",
        "# - name: float64",
        "#   unit: m / s",
        "#   datatype: float64",
        "#   description: descr_float64",
        "#   meta: !!omap",
        "#   - {meta float64: 1}",
        "#   - {a: 2}",
        "# - name: str",
        "#   unit: m / s",
        "#   datatype: string",
        "#   description: descr_str",
        "#   meta: !!omap",
        "#   - {meta str: 1}",
        "#   - {a: 2}",
        "# meta: !!omap",
        "# - comments: [comment1, comment2]",
        "# - {a: 3}",
        "# schema: astropy-2.0",
        "bool int64 float64 str",
        'False 0 0.0 "ab 0"',
        'True 1 1.0 "ab, 1"',
        "False 2 2.0 ab2",
    ]

    out = StringIO()
    t.write(out, format="ascii.ecsv")
    assert out.getvalue().splitlines() == lines


def test_read_float128_pyarrow_fail():
    text = """
# %ECSV 1.0
# ---
# datatype:
# - {name: col0, datatype: float128}
# schema: astropy-2.0
col0
1.5
"""
    with pytest.raises(
        TypeError,
        match="pyarrow engine does not support float128, choose a different engine",
    ):
        Table.read(text, format="ecsv", engine="pyarrow")


def test_write_read_roundtrip(format_engine):
    """
    Write a full-featured table with all types and see that it round-trips on
    readback.  Use both space and comma delimiters.
    """
    t = T_DTYPES.copy()

    # pyarrow does not support float128 and there is no workaround
    if format_engine.get("engine") == "pyarrow" and "float128" in t.colnames:
        del t["float128"]

    for delimiter in DELIMITERS:
        out = StringIO()
        t.write(out, delimiter=delimiter, **patch_format_write(format_engine))

        t2s = [
            Table.read(out.getvalue(), **format_engine),
        ]
        if format_engine["format"] == "ascii.ecsv":
            t2s.extend(
                [
                    Table.read(out.getvalue(), format="ascii"),
                    ascii.read(out.getvalue()),
                    ascii.read(out.getvalue(), format="ecsv", guess=False),
                    ascii.read(out.getvalue(), format="ecsv"),
                ]
            )
        for t2 in t2s:
            assert t.meta == t2.meta
            for name in t.colnames:
                assert t[name].attrs_equal(t2[name])
                assert np.all(t[name] == t2[name])


def test_write_read_roundtrip_empty_table(tmp_path):
    # see https://github.com/astropy/astropy/issues/13191
    sfile = tmp_path / "x.ecsv"
    Table().write(sfile)
    t = Table.read(sfile)
    assert len(t) == 0
    assert len(t.colnames) == 0


def test_bad_delimiter():
    """
    Passing a delimiter other than space or comma gives an exception
    """
    out = StringIO()
    with pytest.raises(ValueError) as err:
        T_DTYPES.write(out, format="ascii.ecsv", delimiter="|")
    assert "only space and comma are allowed" in str(err.value)


@pytest.mark.parametrize("name", ["# name", ' #name " '])
@pytest.mark.parametrize("delimiter", [" ", ","])
def test_stressing_colname_starts_with_hash_etc(format_engine, name, delimiter):
    """Column name starting with # that looks like a comment, see #18710.

    Also names that contain leading/trailing whitespace and a quote character.
    """
    out = io.StringIO()
    t = Table()
    t[name] = [1, 2]
    t["a"] = [3, 4]
    t.write(out, delimiter=delimiter, **patch_format_write(format_engine))
    out.seek(0)
    t2 = Table.read(out.getvalue(), **format_engine)
    assert t2.colnames == [name, "a"]


def test_unavailable_ecsv_engine_for_writing():
    out = io.StringIO()
    t = Table()
    with pytest.raises(
        ValueError,
        match=r"^engine='pyarrow' is not a supported engine for writing, use 'io.ascii'$",
    ):
        t.write(out, format="ecsv", engine="pyarrow")


def test_bad_header_start(format_engine):
    """
    Bad header without initial # %ECSV x.x
    """
    lines = copy.copy(SIMPLE_LINES)
    lines[0] = "# %ECV 0.9"
    kwargs = {"guess": False} if format_engine["format"] == "ascii.ecsv" else {}
    with pytest.raises(ascii.InconsistentTableError):
        Table.read("\n".join(lines), **format_engine, **kwargs)


def test_bad_delimiter_input(format_engine):
    """
    Illegal delimiter in input
    """
    lines = copy.copy(SIMPLE_LINES)
    lines.insert(2, "# delimiter: |")
    kwargs = {"guess": False} if format_engine["format"] == "ascii.ecsv" else {}
    with pytest.raises(ValueError) as err:
        Table.read("\n".join(lines), **format_engine, **kwargs)
    assert "only space and comma are allowed" in str(err.value)


def test_multidim_only_masked(format_engine):
    """Multi-dimensional column with one masked entry

    This hits code in ``get_str_vals`` that requires a pure Python representation of
    data instead of a numpy string array, because later on there can be substitution
    of a longer value into the array element.
    """
    txt = """
# %ECSV 1.0
# ---
# datatype:
# - {name: array3x2, datatype: string, subtype: 'float64[3,2]'}
array3x2
""
"""
    t = Table.read(txt, **format_engine)
    assert len(t) == 1
    assert np.all(t["array3x2"].mask == [[[True, True], [True, True], [True, True]]])


def test_multidim_input(format_engine):
    """
    Multi-dimensional column in input
    """
    t = Table()
    t["a"] = np.arange(24).reshape(2, 3, 4)
    t["a"].info.description = "description"
    t["a"].info.meta = {1: 2}
    t["b"] = [1, 2]

    out = StringIO()
    t.write(out, format="ascii.ecsv")
    t2 = Table.read(out.getvalue(), **format_engine)

    assert np.all(t2["a"] == t["a"])
    assert t2["a"].shape == t["a"].shape
    assert t2["a"].dtype == t["a"].dtype
    assert t2["a"].info.description == t["a"].info.description
    assert t2["a"].info.meta == t["a"].info.meta

    assert np.all(t2["b"] == t["b"])


def test_structured_input(format_engine):
    """
    Structured column in input.
    """
    t = Table()
    # Add unit, description and meta to make sure that round-trips as well.
    t["a"] = Column(
        [("B", (1.0, [2.0, 3.0])), ("A", (9.0, [8.0, 7.0]))],
        dtype=[("s", "U1"), ("v", [("p0", "f8"), ("p1", "2f8")])],
        description="description",
        format=">",  # Most formats do not work with structured!
        unit="m",  # Overall unit should round-trip.
        meta={1: 2},
    )
    t["b"] = Column(
        [[(1.0, 2.0), (9.0, 8.0)], [(3.0, 4.0), (7.0, 6.0)]],
        dtype="f8,f8",
        unit=u.Unit("m,s"),  # Per part unit should round-trip too.
    )

    out = StringIO()
    t.write(out, format="ascii.ecsv")
    t2 = Table.read(out.getvalue(), **format_engine)

    for col in t.colnames:
        assert np.all(t2[col] == t[col])
        assert t2[col].shape == t[col].shape
        assert t2[col].dtype == t[col].dtype
        assert t2[col].unit == t[col].unit
        assert t2[col].format == t[col].format
        assert t2[col].info.description == t[col].info.description
        assert t2[col].info.meta == t[col].info.meta


def test_round_trip_empty_table(format_engine):
    """Test fix in #5010 for issue #5009 (ECSV fails for empty type with bool type)"""
    t = Table(dtype=[bool, "i", "f"], names=["a", "b", "c"])
    out = StringIO()
    t.write(out, format="ascii.ecsv")
    t2 = Table.read(out.getvalue(), **format_engine)
    assert t.dtype == t2.dtype
    assert len(t2) == 0


def test_csv_ecsv_colnames_mismatch(format_engine):
    """
    Test that mismatch in column names from normal CSV header vs.
    ECSV YAML header raises the expected exception.
    """
    lines = copy.copy(SIMPLE_LINES)
    header_index = lines.index("a b c")
    lines[header_index] = "a b d"
    with pytest.raises(ValueError) as err:
        Table.read(lines, **format_engine)
    assert "column names from ECSV header ['a', 'b', 'c']" in str(err.value)


def test_regression_5604():
    """
    See https://github.com/astropy/astropy/issues/5604 for more.
    """
    t = Table()
    t.meta = {"foo": 5 * u.km, "foo2": u.s}
    t["bar"] = [7] * u.km

    out = StringIO()
    t.write(out, format="ascii.ecsv")

    assert "!astropy.units.Unit" in out.getvalue()
    assert "!astropy.units.Quantity" in out.getvalue()


def assert_objects_equal(obj1, obj2, attrs, compare_class=True):
    if compare_class:
        assert obj1.__class__ is obj2.__class__

    assert obj1.shape == obj2.shape

    info_attrs = [
        "info.name",
        "info.format",
        "info.unit",
        "info.description",
        "info.dtype",
    ]
    for attr in attrs + info_attrs:
        a1 = obj1
        a2 = obj2
        for subattr in attr.split("."):
            try:
                a1 = getattr(a1, subattr)
                a2 = getattr(a2, subattr)
            except AttributeError:
                a1 = a1[subattr]
                a2 = a2[subattr]

        if isinstance(a1, np.ndarray) and a1.dtype.kind == "f":
            assert quantity_allclose(a1, a2, rtol=1e-10)
        else:
            assert np.all(a1 == a2)

    # Check meta values and key order but allow None to be equivalent to {}
    meta1 = obj1.info.meta or {}
    meta2 = obj2.info.meta or {}
    assert meta1 == meta2
    assert meta1.keys() == meta2.keys()

    # For no attrs that means we just compare directly.
    if not attrs:
        if isinstance(obj1, np.ndarray) and obj1.dtype.kind == "f":
            assert quantity_allclose(obj1, obj2, rtol=1e-15)
        else:
            assert np.all(obj1 == obj2)


def test_ecsv_mixins_ascii_read_class():
    """Ensure that ascii.read(ecsv_file) returns the correct class
    (QTable if any Quantity subclasses, Table otherwise).
    """
    # Make a table with every mixin type except Quantities
    t = QTable(
        {
            name: col
            for name, col in mixin_cols.items()
            if not isinstance(col.info, QuantityInfo)
        }
    )
    out = StringIO()
    t.write(out, format="ascii.ecsv")
    t2 = ascii.read(out.getvalue(), format="ecsv")
    assert type(t2) is Table

    # Add a single quantity column
    t["lon"] = mixin_cols["lon"]

    out = StringIO()
    t.write(out, format="ascii.ecsv")
    t2 = ascii.read(out.getvalue(), format="ecsv")
    assert type(t2) is QTable


def test_ecsv_mixins_qtable_to_table(format_engine):
    """Test writing as QTable and reading as Table.  Ensure correct classes
    come out.
    """
    names = sorted(mixin_cols)

    t = QTable([mixin_cols[name] for name in names], names=names)
    out = StringIO()
    t.write(out, format="ascii.ecsv")
    t2 = Table.read(out.getvalue(), **format_engine)

    assert t.colnames == t2.colnames

    for name, col in t.columns.items():
        col2 = t2[name]
        attrs = compare_attrs[name]
        compare_class = True

        if isinstance(col.info, QuantityInfo):
            # Downgrade Quantity to Column + unit
            assert type(col2) is Column
            # Class-specific attributes like `value` or `wrap_angle` are lost.
            attrs = ["unit"]
            compare_class = False
            # Compare data values here (assert_objects_equal doesn't know how in this case)
            assert np.allclose(col.value, col2, rtol=1e-10)

        assert_objects_equal(col, col2, attrs, compare_class)


@pytest.mark.parametrize("table_cls", (Table, QTable))
def test_ecsv_mixins_as_one(table_cls, format_engine):
    """Test write/read all cols at once and validate intermediate column names"""
    names = sorted(mixin_cols)
    all_serialized_names = []
    # ECSV stores times as value by default, so we just get the column back.
    # One exception is tm3, which is set to serialize via jd1 and jd2.
    for name in names:
        s_names = serialized_names[name]
        if not name.startswith("tm3"):
            s_names = [
                s_name.replace(".jd1", "")
                for s_name in s_names
                if not s_name.endswith("jd2")
            ]
        all_serialized_names.extend(s_names)

    t = table_cls([mixin_cols[name] for name in names], names=names)

    out = StringIO()
    t.write(out, format="ascii.ecsv")
    t2 = table_cls.read(out.getvalue(), **format_engine)

    assert t.colnames == t2.colnames

    # Read as a ascii.basic table (skip all the ECSV junk)
    t3 = table_cls.read(out.getvalue(), format="ascii.basic")
    assert t3.colnames == all_serialized_names


def make_multidim(col, ndim):
    """Take a col with length=2 and make it N-d by repeating elements.

    For the special case of ndim==1 just return the original.

    The output has shape [3] * ndim. By using 3 we can be sure that repeating
    the two input elements gives an output that is sufficiently unique for
    the multidim tests.
    """
    if ndim > 1:
        idxs = [i % 2 for i in range(3**ndim)]
        col = col[idxs].reshape([3] * ndim)
    return col


@pytest.mark.parametrize("name_col", list(mixin_cols.items()))
@pytest.mark.parametrize("table_cls", (Table, QTable))
@pytest.mark.parametrize("ndim", (1, 2, 3))
def test_ecsv_mixins_per_column(table_cls, name_col, ndim, format_engine):
    """Test write/read one col at a time and do detailed validation.
    This tests every input column type as 1-d, 2-d and 3-d.
    """
    name, col = name_col

    c = make_multidim(np.array([1.0, 2.0]), ndim)
    col = make_multidim(col, ndim)
    t = table_cls([c, col, c], names=["c1", name, "c2"])
    t[name].info.description = "description"
    t[name].info.meta = {"b": 2, "a": 1}

    out = StringIO()
    t.write(out, format="ascii.ecsv")
    t2 = table_cls.read(out.getvalue(), **format_engine)

    assert t.colnames == t2.colnames

    for colname in t.colnames:
        assert len(t2[colname].shape) == ndim
        if colname in ("c1", "c2"):
            compare = ["data"]
        else:
            # Storing Longitude as Column loses wrap_angle.
            compare = [
                attr
                for attr in compare_attrs[colname]
                if not (attr == "wrap_angle" and table_cls is Table)
            ]
        assert_objects_equal(t[colname], t2[colname], compare)

    # Special case to make sure Column type doesn't leak into Time class data
    if name.startswith("tm"):
        assert t2[name]._time.jd1.__class__ is np.ndarray
        assert t2[name]._time.jd2.__class__ is np.ndarray


def test_round_trip_masked_table_default(tmp_path):
    """Test (mostly) round-trip of MaskedColumn through ECSV using default serialization
    that uses an empty string "" to mark NULL values.  Note:

    >>> simple_table(masked=True)
    <Table masked=True length=3>
      a      b     c
    int64 float64 str1
    ----- ------- ----
       --     1.0    c
        2     2.0   --
        3      --    e
    """
    filename = tmp_path / "test.ecsv"

    t = simple_table(masked=True)  # int, float, and str cols with one masked element
    t.write(filename)

    t2 = Table.read(filename)
    assert t2.masked is False
    assert t2.colnames == t.colnames
    for name in t2.colnames:
        # From formal perspective the round-trip columns are the "same"
        assert np.all(t2[name].mask == t[name].mask)
        assert np.all(t2[name] == t[name])

        # But peeking under the mask shows that the underlying data are changed
        # because by default ECSV uses "" to represent masked elements.
        t[name].mask = False
        t2[name].mask = False
        assert not np.all(t2[name] == t[name])  # Expected diff


def test_round_trip_masked_table_serialize_mask(tmp_path, format_engine):
    """
    Same as prev but set the serialize_method to 'data_mask' so mask is written out
    """
    filename = tmp_path / "test.ecsv"

    t = simple_table(masked=True)  # int, float, and str cols with one masked element
    t["c"][0] = ""  # This would come back as masked for default "" NULL marker

    # MaskedColumn with no masked elements. See table the MaskedColumnInfo class
    # _represent_as_dict() method for info about how we test a column with no masked elements.
    t["d"] = [1, 2, 3]

    t.write(filename, serialize_method="data_mask")

    t2 = Table.read(filename, **format_engine)
    assert t2.masked is False
    assert t2.colnames == t.colnames
    for name in t2.colnames:
        assert np.all(t2[name].mask == t[name].mask)
        assert np.all(t2[name] == t[name])

        # Data under the mask round-trips also (unmask data to show this).
        t[name].mask = False
        t2[name].mask = False
        assert np.all(t2[name] == t[name])


@pytest.mark.parametrize("table_cls", (Table, QTable))
def test_ecsv_round_trip_user_defined_unit(table_cls, tmp_path, format_engine):
    """Ensure that we can read-back enabled user-defined units."""

    # Test adapted from #8897, where it was noted that this works
    # but was not tested.
    filename = tmp_path / "test.ecsv"
    unit = u.def_unit("bandpass_sol_lum")
    t = table_cls()
    t["l"] = np.arange(5) * unit
    t.write(filename)
    # without the unit enabled, get UnrecognizedUnit
    if table_cls is QTable:
        ctx = pytest.warns(u.UnitsWarning, match=r"'bandpass_sol_lum' did not parse .*")
    else:
        ctx = nullcontext()
    # Note: The read might also generate ResourceWarning, in addition to UnitsWarning
    with ctx:
        t2 = table_cls.read(filename, **format_engine)
    assert isinstance(t2["l"].unit, u.UnrecognizedUnit)
    assert str(t2["l"].unit) == "bandpass_sol_lum"
    if table_cls is QTable:
        assert np.all(t2["l"].value == t["l"].value)
    else:
        assert np.all(t2["l"] == t["l"])

    # But with it enabled, it works.
    with u.add_enabled_units(unit):
        t3 = table_cls.read(filename)
        assert t3["l"].unit is unit
        assert np.all(t3["l"] == t["l"])

        # Just to be sure, also try writing with unit enabled.
        filename2 = tmp_path / "test2.ecsv"
        t3.write(filename2)
        t4 = table_cls.read(filename, **format_engine)
        assert t4["l"].unit is unit
        assert np.all(t4["l"] == t["l"])


def test_read_masked_bool(format_engine):
    txt = """\
# %ECSV 1.0
# ---
# datatype:
# - {name: col0, datatype: bool}
# schema: astropy-2.0
col0
1
0
True
""
False
"""
    dat = Table.read(txt, **format_engine)
    col = dat["col0"]
    assert isinstance(col, MaskedColumn)
    assert np.all(col.mask == [False, False, False, True, False])
    assert np.all(col == [True, False, True, False, False])


@pytest.mark.parametrize("serialize_method", ["null_value", "data_mask"])
@pytest.mark.parametrize("dtype", [np.int64, np.float64, bool, str])
@pytest.mark.parametrize("delimiter", [",", " "])
def test_roundtrip_multidim_masked_array(
    serialize_method, dtype, delimiter, format_engine
):
    # TODO also test empty string with null value
    t = Table()
    col = MaskedColumn(np.arange(12).reshape(2, 3, 2), dtype=dtype)
    if dtype is str:
        # np does something funny and gives a dtype of U21.
        col = col.astype("U2")
    col.mask[0, 0, 0] = True
    col.mask[1, 1, 1] = True
    t["a"] = col
    t["b"] = ["x", "y"]  # Add another column for kicks
    out = StringIO()
    t.write(
        out, format="ascii.ecsv", serialize_method=serialize_method, delimiter=delimiter
    )
    t2 = Table.read(out.getvalue(), **format_engine)

    assert t2.masked is False
    assert t2.colnames == t.colnames
    for name in t2.colnames:
        assert t2[name].dtype == t[name].dtype
        if hasattr(t[name], "mask"):
            assert np.all(t2[name].mask == t[name].mask)
        assert np.all(t2[name] == t[name])


@pytest.mark.parametrize("subtype", ["some-user-type", "complex"])
def test_multidim_unknown_subtype(subtype):
    """Test an ECSV file with a string type but unknown subtype"""
    txt = f"""\
# %ECSV 1.0
# ---
# datatype:
# - name: a
#   datatype: string
#   subtype: {subtype}
# schema: astropy-2.0
a
[1,2]
[3,4]"""
    with pytest.warns(
        InvalidEcsvDatatypeWarning,
        match=rf"unexpected subtype '{subtype}' set for column 'a'",
    ):
        t = Table.read(txt, format="ascii.ecsv")

    assert t["a"].dtype.kind == "U"
    assert t["a"][0] == "[1,2]"


def test_multidim_bad_shape(format_engine):
    """Test a malformed ECSV file"""
    txt = """\
# %ECSV 1.0
# ---
# datatype:
# - name: a
#   datatype: string
#   subtype: int64[3]
# schema: astropy-2.0
a
[1,2]
[3,4]"""
    with pytest.raises(
        ValueError, match="column 'a' failed to convert: shape mismatch"
    ):
        Table.read(txt, **format_engine)


def test_write_not_json_serializable():
    t = Table()
    t["a"] = np.array([{1, 2}, 1], dtype=object)
    match = (
        "could not convert column 'a' to string: Object of type set is not JSON"
        " serializable"
    )
    out = StringIO()
    with pytest.raises(TypeError, match=match):
        t.write(out, format="ascii.ecsv")


def test_read_not_json_serializable(format_engine):
    """Test a malformed ECSV file"""
    txt = """\
# %ECSV 1.0
# ---
# datatype:
# - {name: a, datatype: string, subtype: json}
# schema: astropy-2.0
a
fail
[3,4]"""
    match = "column 'a' failed to convert: column value is not valid JSON"
    with pytest.raises(ValueError, match=match):
        Table.read(txt, **format_engine)


def test_read_bad_datatype(format_engine):
    """Test a malformed ECSV file"""
    txt = """\
# %ECSV 1.0
# ---
# datatype:
# - {name: a, datatype: object}
# schema: astropy-2.0
a
{"x":1}
[3,4]"""
    with pytest.warns(
        InvalidEcsvDatatypeWarning,
        match="unexpected datatype 'object' of column 'a' is not in allowed",
    ):
        t = Table.read(txt, **format_engine)
    col = t["a"]
    if format_engine["format"] == "ascii.ecsv":
        # This is the legacy "passing" behavior for io.ascii.ecsv, but it is just wrong.
        # Not clear why this was included as expected behavior.
        assert col[0] == '{"x":1}'
        assert col[1] == "[3,4]"
        assert all(type(col[ii]) is str for ii in range(len(col)))
        assert col.dtype.kind == "O"
    else:
        assert col[0] == {"x": 1}
        assert type(col[0]) is dict
        assert col[1] == [3, 4]
        assert type(col[1]) is list
        assert col.dtype.kind == "O"


def test_read_complex(format_engine):
    """Test an ECSV v1.0 file with a complex column"""
    txt = """\
# %ECSV 1.0
# ---
# datatype:
# - {name: a, datatype: complex}
# schema: astropy-2.0
a
1+1j
2+2j"""

    with pytest.warns(
        InvalidEcsvDatatypeWarning,
        match="unexpected datatype 'complex' of column 'a' is not in allowed",
    ):
        t = Table.read(txt, **format_engine)
    assert t["a"].dtype.type is np.complex128


def test_read_str(format_engine):
    """Test an ECSV file with a 'str' instead of 'string' datatype"""
    txt = """\
# %ECSV 1.0
# ---
# datatype:
# - {name: a, datatype: str}
# schema: astropy-2.0
a
sometext
S"""  # also testing single character text

    with pytest.warns(
        InvalidEcsvDatatypeWarning,
        match="unexpected datatype 'str' of column 'a' is not in allowed",
    ):
        t = Table.read(txt, **format_engine)
    assert isinstance(t["a"][1], str)
    assert isinstance(t["a"][0], np.str_)


def test_read_bad_datatype_for_object_subtype(format_engine):
    """Test a malformed ECSV file"""
    txt = """\
# %ECSV 1.0
# ---
# datatype:
# - {name: a, datatype: int64, subtype: json}
# schema: astropy-2.0
a
fail
[3,4]"""
    match = "column 'a' failed to convert: datatype of column 'a' must be \"string\""
    with pytest.raises(ValueError, match=match):
        Table.read(txt, **format_engine)


def test_full_repr_roundtrip(format_engine):
    """Test round-trip of float values to full precision even with format
    specified"""
    t = Table()
    t["a"] = np.array([np.pi, 1 / 7], dtype=np.float64)
    t["a"].info.format = ".2f"
    out = StringIO()
    t.write(out, format="ascii.ecsv")
    t2 = Table.read(out.getvalue(), **format_engine)
    # Very small differences in float64 values with pandas. TODO: check why.
    atol = 1e-16 if format_engine.get("engine") == "pandas" else 0.0
    np.testing.assert_allclose(t["a"], t2["a"], atol=atol, rtol=0)
    assert t2["a"].info.format == ".2f"


#############################################################################
# Define a number of specialized columns for testing and the expected values
# of `datatype` for each column.
#############################################################################


# First here is some helper code used to make the expected outputs code.
def _get_ecsv_header_dict(text):
    lines = [line.strip() for line in text.splitlines()]
    lines = [line[2:] for line in lines if line.startswith("#")]
    lines = lines[2:]  # Get rid of the header
    return yaml.safe_load("\n".join(lines))


def _make_expected_values(cols):
    for name, col in cols.items():
        t = Table()
        t[name] = col
        out = StringIO()
        t.write(out, format="ascii.ecsv")
        hdr = _get_ecsv_header_dict(out.getvalue())
        fmt_hdr = pformat(hdr["datatype"])
        print(f"exps[{name!r}] =", fmt_hdr[:1])
        print(fmt_hdr[1:])
        print()


# Expected values of `datatype` for each column
exps = {}
cols = {}

# Run of the mill scalar for completeness
cols["scalar"] = np.array([1, 2], dtype=np.int16)
exps["scalar"] = [{"datatype": "int16", "name": "scalar"}]

# Array of lists that works as a 2-d variable array. This is just treated
# as an object.
cols["2-d variable array lists"] = c = np.empty(shape=(2,), dtype=object)
c[0] = [[1, 2], ["a", 4]]
c[1] = [[1, 2, 3], [4, 5.25, 6]]
exps["2-d variable array lists"] = [
    {"datatype": "string", "name": "2-d variable array lists", "subtype": "json"}
]

# Array of numpy arrays that is a 2-d variable array
cols["2-d variable array numpy"] = c = np.empty(shape=(2,), dtype=object)
c[0] = np.array([[1, 2], [3, 4]], dtype=np.float32)
c[1] = np.array([[1, 2, 3], [4, 5.5, 6]], dtype=np.float32)
exps["2-d variable array numpy"] = [
    {
        "datatype": "string",
        "name": "2-d variable array numpy",
        "subtype": "float32[2,null]",
    }
]

cols["1-d variable array lists"] = np.array([[1, 2], [3, 4, 5]], dtype=object)
exps["1-d variable array lists"] = [
    {"datatype": "string", "name": "1-d variable array lists", "subtype": "json"}
]

# Variable-length array
cols["1-d variable array numpy"] = np.array(
    [np.array([1, 2], dtype=np.uint8), np.array([3, 4, 5], dtype=np.uint8)],
    dtype=object,
)
exps["1-d variable array numpy"] = [
    {"datatype": "string", "name": "1-d variable array numpy", "subtype": "uint8[null]"}
]

cols["1-d variable array numpy str"] = np.array(
    [np.array(["a", "b"]), np.array(["c", "d", "e"])], dtype=object
)
exps["1-d variable array numpy str"] = [
    {
        "datatype": "string",
        "name": "1-d variable array numpy str",
        "subtype": "string[null]",
    }
]

cols["1-d variable array numpy bool"] = np.array(
    [np.array([True, False]), np.array([True, False, True])], dtype=object
)
exps["1-d variable array numpy bool"] = [
    {
        "datatype": "string",
        "name": "1-d variable array numpy bool",
        "subtype": "bool[null]",
    }
]

cols["1-d regular array"] = np.array([[1, 2], [3, 4]], dtype=np.int8)
exps["1-d regular array"] = [
    {"datatype": "string", "name": "1-d regular array", "subtype": "int8[2]"}
]

cols["2-d regular array"] = np.arange(8, dtype=np.float16).reshape(2, 2, 2)
exps["2-d regular array"] = [
    {"datatype": "string", "name": "2-d regular array", "subtype": "float16[2,2]"}
]

cols["scalar object"] = np.array([{"a": 1}, {"b": 2}], dtype=object)
exps["scalar object"] = [
    {"datatype": "string", "name": "scalar object", "subtype": "json"}
]

cols["1-d object"] = np.array(
    [[{"a": 1}, {"b": 2}], [{"a": 1}, {"b": 2}]], dtype=object
)
exps["1-d object"] = [
    {"datatype": "string", "name": "1-d object", "subtype": "json[2]"}
]


@pytest.mark.parametrize("name,col,exp", list(zip(cols, cols.values(), exps.values())))
def test_specialized_columns(name, col, exp, format_engine):
    """Test variable length lists, multidim columns, object columns."""
    t = Table()
    t[name] = col
    out = StringIO()
    t.write(out, format="ascii.ecsv")
    hdr = _get_ecsv_header_dict(out.getvalue())
    assert hdr["datatype"] == exp
    t2 = Table.read(out.getvalue(), **format_engine)
    assert t2.colnames == t.colnames
    for colname in t2.colnames:
        assert t2[colname].dtype == t[colname].dtype
        for val1, val2 in zip(t2[colname], t[colname]):
            if isinstance(val1, np.ndarray):
                assert val1.dtype == val2.dtype
            assert np.all(val1 == val2)


def test_full_subtypes(format_engine):
    """Read ECSV file created by M. Taylor that includes scalar, fixed array,
    variable array for all datatypes. This file has missing values for all
    columns as both per-value null and blank entries for the entire column
    value.

    Note: original file was modified to include blank values in f_float and
    f_double columns.
    """
    t = Table.read(os.path.join(TEST_DIR, "data", "subtypes.ecsv"), **format_engine)
    colnames = [
        "i_index",
        "s_byte",
        "s_short",
        "s_int",
        "s_long",
        "s_float",
        "s_double",
        "s_string",
        "s_boolean",
        "f_byte",
        "f_short",
        "f_int",
        "f_long",
        "f_float",
        "f_double",
        "f_string",
        "f_boolean",
        "v_byte",
        "v_short",
        "v_int",
        "v_long",
        "v_float",
        "v_double",
        "v_string",
        "v_boolean",
        "m_int",
        "m_double",
    ]
    assert t.colnames == colnames

    type_map = {
        "byte": "int8",
        "short": "int16",
        "int": "int32",
        "long": "int64",
        "float": "float32",
        "double": "float64",
        "string": "str",
        "boolean": "bool",
    }

    for col in t.itercols():
        info = col.info
        if info.name == "i_index":
            continue

        assert isinstance(col, MaskedColumn)

        type_name = info.name[2:]  # short, int, etc
        subtype = info.name[:1]

        if subtype == "s":  # Scalar
            assert col.shape == (16,)

        if subtype == "f":  # Fixed array
            assert col.shape == (16, 3)

        if subtype == "v":  # Variable array
            assert col.shape == (16,)
            assert info.dtype.name == "object"
            for val in col:
                assert isinstance(val, np.ndarray)
                assert val.dtype.name.startswith(type_map[type_name])
                assert len(val) in [0, 1, 2, 3]
        else:
            assert info.dtype.name.startswith(type_map[type_name])


def test_masked_empty_subtypes(format_engine):
    """Test blank field in subtypes. Similar to previous test but with explicit
    checks of values"""
    txt = """
# %ECSV 1.0
# ---
# datatype:
# - {name: o, datatype: string, subtype: json}
# - {name: f, datatype: string, subtype: 'int64[2]'}
# - {name: v, datatype: string, subtype: 'int64[null]'}
# schema: astropy-2.0
o f v
null [0,1] [1]
"" "" ""
[1,2] [2,3] [2,3]
"""
    t = Table.read(txt, **format_engine)
    assert np.all(t["o"] == np.array([None, -1, [1, 2]], dtype=object))
    assert np.all(t["o"].mask == [False, True, False])

    exp = np.ma.array([[0, 1], [-1, -1], [2, 3]], mask=[[0, 0], [1, 1], [0, 0]])
    assert np.all(t["f"] == exp)
    assert np.all(t["f"].mask == exp.mask)

    assert np.all(t["v"][0] == [1])
    assert np.all(t["v"][2] == [2, 3])
    assert np.all(t["v"].mask == [False, True, False])


def test_masked_vals_in_array_subtypes(format_engine):
    """Test null values in fixed and variable array subtypes."""
    t = Table()
    t["f"] = np.ma.array([[1, 2], [3, 4]], mask=[[0, 1], [1, 0]], dtype=np.int64)
    t["v"] = np.empty(2, dtype=object)
    t["v"][0] = np.ma.array([1, 2], mask=[0, 1], dtype=np.int64)
    t["v"][1] = np.ma.array([3, 4, 5], mask=[1, 0, 0], dtype=np.int64)

    out = StringIO()
    t.write(out, format="ascii.ecsv")
    txt = """
    # %ECSV 1.0
    # ---
    # datatype:
    # - {name: f, datatype: string, subtype: 'int64[2]'}
    # - {name: v, datatype: string, subtype: 'int64[null]'}
    # schema: astropy-2.0
    f v
    [1,null] [1,null]
    [null,4] [null,4,5]
    """
    hdr = _get_ecsv_header_dict(out.getvalue())
    hdr_exp = _get_ecsv_header_dict(txt)
    assert hdr == hdr_exp
    t2 = Table.read(out.getvalue(), **format_engine)
    assert t2.colnames == t.colnames
    for name in t2.colnames:
        assert t2[name].dtype == t[name].dtype
        assert type(t2[name]) is type(t[name])
        for val1, val2 in zip(t2[name], t[name]):
            if isinstance(val1, np.ndarray):
                assert val1.dtype == val2.dtype
            if isinstance(val1, np.ma.MaskedArray):
                assert np.all(val1.mask == val2.mask)
            assert np.all(val1 == val2)


def test_guess_ecsv_with_one_column():
    """Except for ECSV, guessing always requires at least 2 columns"""
    txt = """
    # %ECSV 1.0
    # ---
    # datatype:
    # - {name: col, datatype: string, description: hello}
    # schema: astropy-2.0
    col
    1
    2
    """
    t = ascii.read(txt)
    assert t["col"].dtype.kind == "U"  # would be int with basic format
    assert t["col"].description == "hello"


@pytest.mark.parametrize("masked", [MaskedColumn, Masked, np.ma.MaskedArray])
def test_write_structured_masked_column(masked, format_engine):
    a = np.array([(1, 2), (3, 4)], dtype="i,i")
    mc = masked(a, mask=[(True, False), (False, False)])
    t = Table([mc], names=["mc"])
    out = StringIO()
    t.write(out, format="ascii.ecsv")
    t2 = Table.read(out.getvalue(), **format_engine)
    assert type(t2["mc"]) is type(t["mc"])
    assert (t2["mc"] == mc).all()
    assert (t2["mc"].mask == mc.mask).all()


def test_write_masked_time_ymdhms_mixin(format_engine):
    # Regression test for gh-16370
    # Make a masked time,
    t = Time({"year": 2000, "month": 1, "day": [1, 2]})
    t[0] = np.ma.masked
    # Create a table and write to a file
    qt = QTable([t], names=["t"])
    out = StringIO()
    qt.write(out, format="ascii.ecsv")
    # Read back and compare.
    qt2 = QTable.read(out.getvalue(), **format_engine)
    # Note that value under time does not roundtrip
    assert (qt2["t"] == t).all()
    assert (qt2["t"].mask == t.mask).all()


def test_register_bad_engine():
    msg = "Subclasses of ECSVEngine must define a class attribute 'name' as a string, got <class 'int'>."
    with pytest.raises(TypeError, match=msg):

        class BadEngine(ECSVEngine):
            name = 1
            format = "ascii.ecsv"


@pytest.mark.parametrize(
    "compressed_filename",
    [
        "test.ecsv.gz",
        pytest.param(
            "test.ecsv.bz2",
            marks=pytest.mark.xfail(not HAS_BZ2, reason="no bz2 support"),
        ),
        pytest.param(
            "test.ecsv.xz",
            marks=pytest.mark.xfail(not HAS_LZMA, reason="no lzma support"),
        ),
    ],
)
def test_compressed_files(tmp_path, format_engine, compressed_filename):
    filename = tmp_path / "test.ecsv"
    t = simple_table()
    t.write(filename, format="ascii.ecsv")

    compressed_filename = tmp_path / compressed_filename
    if compressed_filename.suffix == ".gz":
        import gzip

        opener = gzip.open
    elif compressed_filename.suffix == ".bz2":
        import bz2

        opener = bz2.open
    elif compressed_filename.suffix == ".xz":
        import lzma

        opener = lzma.open
    else:
        # Shouldn't really happen
        opener = open

    # Compress ecsv file
    with open(filename, "rb") as f_in:
        with opener(compressed_filename, "wb") as f_out:
            f_out.writelines(f_in)

    # Open compressed file and compare to ensure it's read correctly
    t_comp = Table.read(compressed_filename, **format_engine)
    assert_array_equal(t, t_comp)
