"""Contracts of the agent-reviewed strategy: without score files it holds
exactly what its base strategy holds, a scores file only steers sessions it
was stamped before, and a vetoed holding is replaced the same session."""

import json
import os
import tempfile
import unittest

import numpy as np
import pandas as pd

from strategies.agent_review import AgentReviewStrategy, session_deadline
from strategies.robust_momentum import RobustMomentumStrategy
from tests.test_robust_strategies import _bars


def _write(score_dir: str, as_of, created_at, scores: dict) -> None:
    rows = {t: {"score": s, "veto": s == "veto", "summary": "x"} for t, s in scores.items()}
    for v in rows.values():
        if v["veto"]:
            v["score"] = 0
    path = os.path.join(score_dir, f"scores_{pd.Timestamp(as_of).date()}_{len(os.listdir(score_dir)):06d}.json")
    with open(path, "w") as fh:
        json.dump({"as_of": str(pd.Timestamp(as_of).date()), "created_at": pd.Timestamp(created_at).isoformat(),
                   "scores": rows}, fh)


class AgentReviewTests(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.dir = self.tmp.name
        self.bars = _bars(n_days=420, n_assets=40)
        self.kw = dict(top_n=5, pool=20, rebal_freq=5, buffer=0.5, warmup=260)

    def tearDown(self):
        self.tmp.cleanup()

    def _agent(self):
        return AgentReviewStrategy(score_dir=self.dir, **self.kw).generate_signals(self.bars)

    def test_without_score_files_it_holds_what_the_base_strategy_holds(self):
        base = RobustMomentumStrategy("overnight_recovery", top_n=5, rebal_freq=5, buffer=0.5,
                                      warmup=260).generate_signals(self.bars)
        pd.testing.assert_frame_equal(self._agent(), base, atol=1e-12)

    def test_scores_stamped_after_the_session_close_do_not_steer_it(self):
        dates = self.bars["close"].index
        k = 260 + 5 * 20                       # a rebalance row; its weights trade at dates[k + 1]
        before = self._agent()
        worst = before.columns[before.iloc[k + 1] == 0][0]
        late = session_deadline(dates[k + 1]) + pd.Timedelta(minutes=1)
        _write(self.dir, dates[k], late, {worst: 2, **{t: -2 for t in before.columns[before.iloc[k + 1] > 0]}})
        pd.testing.assert_frame_equal(self._agent(), before, atol=1e-12)
        _write(self.dir, dates[k], session_deadline(dates[k + 1]) - pd.Timedelta(hours=1),
               {t: -2 for t in before.columns[before.iloc[k + 1] > 0]})
        after = self._agent()
        pd.testing.assert_frame_equal(after.iloc[:k + 1], before.iloc[:k + 1], atol=1e-12)
        self.assertNotEqual(set(after.columns[after.iloc[k + 1] > 0]), set(before.columns[before.iloc[k + 1] > 0]))

    def test_a_vetoed_holding_is_replaced_between_rebalances(self):
        dates = self.bars["close"].index
        k = 260 + 5 * 20 + 2                   # not a rebalance row
        before = self._agent()
        held = list(before.columns[before.iloc[k + 1] > 0])
        _write(self.dir, dates[k], session_deadline(dates[k + 1]) - pd.Timedelta(hours=2), {held[0]: "veto"})
        after = self._agent()
        now = set(after.columns[after.iloc[k + 1] > 0])
        self.assertNotIn(held[0], now)
        self.assertEqual(len(now), len(held))
        self.assertAlmostEqual(float(after.iloc[k + 1].sum()), 1.0, places=9)
        self.assertTrue(np.allclose(after.iloc[:k + 1], before.iloc[:k + 1]))


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