import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

import polars as pl

from intraday_fragility.schema import integrity_flags


def make_df(rows):
    return pl.DataFrame(
        rows,
        schema={"timestamp": pl.Datetime("ns", "UTC"), "open": pl.Float64,
                "high": pl.Float64, "low": pl.Float64, "close": pl.Float64,
                "volume": pl.Float64, "ticker": pl.String},
        orient="row",
    )


TS = "2020-01-02T14:30:00"


def test_valid_bar_not_flagged():
    import datetime as dt
    ts = dt.datetime(2020, 1, 2, 14, 30, tzinfo=dt.timezone.utc)
    df = make_df([[ts, 10.0, 11.0, 9.5, 10.5, 100.0, "AAA"]])
    f = integrity_flags(df)
    assert not f["flag_ohlc_violation"][0]
    assert not f["flag_nonpositive_price"][0]
    assert not f["flag_negative_volume"][0]
    assert not f["flag_duplicate_key"][0]


def test_high_below_close_flagged():
    import datetime as dt
    ts = dt.datetime(2020, 1, 2, 14, 30, tzinfo=dt.timezone.utc)
    df = make_df([[ts, 10.0, 10.2, 9.5, 10.5, 100.0, "AAA"]])
    assert integrity_flags(df)["flag_ohlc_violation"][0]


def test_low_above_open_flagged():
    import datetime as dt
    ts = dt.datetime(2020, 1, 2, 14, 30, tzinfo=dt.timezone.utc)
    df = make_df([[ts, 9.0, 11.0, 9.5, 10.5, 100.0, "AAA"]])
    assert integrity_flags(df)["flag_ohlc_violation"][0]


def test_duplicate_key_flagged():
    import datetime as dt
    ts = dt.datetime(2020, 1, 2, 14, 30, tzinfo=dt.timezone.utc)
    df = make_df([
        [ts, 10.0, 11.0, 9.5, 10.5, 100.0, "AAA"],
        [ts, 10.0, 11.0, 9.5, 10.5, 200.0, "AAA"],
        [ts, 10.0, 11.0, 9.5, 10.5, 100.0, "BBB"],
    ])
    f = integrity_flags(df)
    assert f["flag_duplicate_key"].sum() == 2
    assert not f["flag_duplicate_key"][2]