from io import StringIO

import locan.constants
from locan.locan_io import load_thunderstorm_file
from locan.locan_io.locdata.thunderstorm_io import load_thunderstorm_header


def test_get_correct_column_names_from_Thunderstorm_header():
    columns = load_thunderstorm_header(
        path=locan.ROOT_DIR / "tests/test_data/Thunderstorm_dstorm_data.csv"
    )
    assert columns == [
        "original_index",
        "frame",
        "position_x",
        "position_y",
        "psf_sigma",
        "intensity",
        "local_background",
        "local_background_sigma",
        "chi_square",
        "uncertainty",
    ]

    file_like = StringIO(
        "id,frame,x [nm],y [nm]\n73897.0,2001.0,1320.109670647555,26344.7124618434"
    )
    columns = load_thunderstorm_header(path=file_like)
    assert columns == ["original_index", "frame", "position_x", "position_y"]


def test_loading_Thunderstorm_file():
    dat = load_thunderstorm_file(
        path=locan.ROOT_DIR / "tests/test_data/Thunderstorm_dstorm_data.csv", nrows=10
    )
    # print(dat.data.columns)
    assert len(dat) == 10

    file_like = StringIO(
        "id,frame,x [nm],y [nm]\n73897.0,2001.0,1320.109670647555,26344.7124618434"
    )
    dat = load_thunderstorm_file(path=file_like, nrows=1)
    assert len(dat) == 1
