from __future__ import annotations

import sys
import unittest
from pathlib import Path

import pandas as pd

_CHAN = Path(__file__).resolve().parents[2]
if str(_CHAN) not in sys.path:
	sys.path.insert(0, str(_CHAN))

from chanlun.pipeline.timeframe import TF_DF  # noqa: E402


def _zigzag_df(n=160, step=8):
	dates = pd.date_range("2024-01-01", periods=n, freq="5min")
	rows = []
	price = 100.0
	for i, date in enumerate(dates):
		up = (i // step) % 2 == 0
		if up:
			o = price
			c = price + 1.5
			h = c + 0.3
			l = o - 0.2
		else:
			o = price
			c = price - 1.5
			h = o + 0.2
			l = c - 0.3
		price = c
		rows.append(
			{
				"date": date,
				"open": o,
				"high": h,
				"low": l,
				"close": c,
				"volume": 1.0,
			}
		)
	return pd.DataFrame(rows)


def _sure_bi_key(bi):
	return (str(bi.start_time), bi.dir.name, round(float(bi.high), 6), round(float(bi.low), 6))


def _zs_key(zs):
	return (
		str(zs.start_time),
		round(float(zs.zg), 6),
		round(float(zs.zd), 6),
		len(zs.bi_list),
	)


class TestIncremental(unittest.TestCase):
	def test_init_stream_matches_batch_push(self):
		df = _zigzag_df()
		stream = TF_DF()
		stream.init_stream(df, 1, "5m")

		batch = TF_DF()
		indexed = batch.add_indicators(df.copy())
		klu = batch.cal_kl_data(indexed)
		klc = []
		last = None
		for k in klu:
			batch._push_klu_into_klc_list(klc, k, last)
			last = k
		batch.klc_list = klc
		batch.rebuild_bi_zs()

		self.assertEqual(len(stream.klu_list), len(klu))
		self.assertEqual(len(stream.klc_list), len(klc))
		self.assertEqual(
			[_sure_bi_key(b) for b in stream.bi_list if b.is_sure],
			[_sure_bi_key(b) for b in batch.bi_list if b.is_sure],
		)
		self.assertEqual(
			[_zs_key(z) for z in stream.bi_zs_list],
			[_zs_key(z) for z in batch.bi_zs_list],
		)

	def test_append_bar_matches_init_stream(self):
		df = _zigzag_df()
		stream = TF_DF()
		stream.init_stream(df, 1, "5m")

		inc = TF_DF()
		for _, row in df.iterrows():
			inc.append_bar(row)

		self.assertEqual(len(inc.klu_list), len(stream.klu_list))
		self.assertEqual(len(inc.klc_list), len(stream.klc_list))
		self.assertEqual(
			[_sure_bi_key(b) for b in inc.bi_list if b.is_sure],
			[_sure_bi_key(b) for b in stream.bi_list if b.is_sure],
		)
		self.assertEqual(
			[_zs_key(z) for z in inc.bi_zs_list],
			[_zs_key(z) for z in stream.bi_zs_list],
		)

	def test_replace_last_bar_keeps_count(self):
		df = _zigzag_df(n=80)
		tf = TF_DF()
		tf.init_stream(df, 1, "5m")
		n_klu = len(tf.klu_list)
		last = df.iloc[-1].copy()
		last["close"] = float(last["close"]) + 0.01
		last["high"] = max(float(last["high"]), float(last["close"]))
		tf.replace_last_bar(last)
		self.assertEqual(len(tf.klu_list), n_klu)
		self.assertGreater(len(tf.klc_list), 0)

	def test_check_fx_skips_forming_right_wing(self):
		from types import SimpleNamespace

		from chanlun.core.ChanEnum import Chan_FX_TYPE

		tf = TF_DF()
		pre = SimpleNamespace(high=10, low=8)
		nxt_open = SimpleNamespace(high=11, low=7, end_klu=None)
		nxt_done = SimpleNamespace(high=11, low=7, end_klu=object())
		center = SimpleNamespace(
			pre=pre,
			next=nxt_open,
			high=12,
			low=9,
			set_fx=lambda *_a, **_k: None,
		)
		self.assertEqual(tf.check_fx(center), Chan_FX_TYPE.UNKNOWN)
		center.next = nxt_done
		self.assertEqual(tf.check_fx(center), Chan_FX_TYPE.TOP)


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