import unittest
from run_study import shift_day_profile

class ShiftTests(unittest.TestCase):
    def setUp(self):
        self.ts = [1_700_000_000_000 + i * 3_600_000 for i in range(6)]
        self.load = [100.0] * 6

    def test_energy_conservation(self):
        original = [90, 80, 70, 110, 130, 120]
        shifted, moved, _, _, _ = shift_day_profile(original, self.load, self.ts, .10, 4)
        self.assertAlmostEqual(sum(original), sum(shifted), places=8)
        self.assertGreaterEqual(moved, 0)

    def test_flat_profile_unchanged(self):
        original = [100.0] * 6
        shifted, moved, _, _, _ = shift_day_profile(original, self.load, self.ts, .10, 4)
        self.assertEqual(original, shifted)
        self.assertEqual(moved, 0)

    def test_source_and_destination_caps(self):
        original = [150, 120, 80, 60, 110, 130]
        shifted, moved, out, incoming, _ = shift_day_profile(original, self.load, self.ts, .05, 6)
        self.assertTrue(all(x <= 5.0 + 1e-8 for x in out))
        self.assertTrue(all(x <= 5.0 + 1e-8 for x in incoming))
        self.assertTrue(all(a + b <= 5.0 + 1e-8 for a, b in zip(out, incoming)))
        self.assertAlmostEqual(sum(out), moved, places=8)
        self.assertAlmostEqual(sum(incoming), moved, places=8)

    def test_window_respected_by_elapsed_time(self):
        stamps = [1_700_000_000_000, 1_700_000_000_000 + 5 * 3_600_000]
        shifted, moved, _, _, max_used = shift_day_profile([150, 50], [100, 100], stamps, .10, 2)
        self.assertEqual(shifted, [150, 50])
        self.assertEqual(moved, 0)
        self.assertEqual(max_used, 0)
        shifted5, moved5, _, _, max_used5 = shift_day_profile([150, 50], [100, 100], stamps, .10, 5)
        self.assertAlmostEqual(sum(shifted5), 200)
        self.assertGreater(moved5, 0)
        self.assertEqual(max_used5, 5)

    def test_no_flex_means_no_change(self):
        original = [150, 120, 80, 60, 110, 130]
        shifted, moved, _, _, _ = shift_day_profile(original, self.load, self.ts, 0, 4)
        self.assertEqual(original, shifted)
        self.assertEqual(moved, 0)

    def test_input_not_mutated(self):
        original = [150, 120, 80, 60, 110, 130]
        copy = list(original)
        shift_day_profile(original, self.load, self.ts, .10, 4)
        self.assertEqual(original, copy)

    def test_invalid_fraction_rejected(self):
        with self.assertRaises(ValueError):
            shift_day_profile([1]*6, self.load, self.ts, .75, 4)

    def test_invalid_window_rejected(self):
        with self.assertRaises(ValueError):
            shift_day_profile([1]*6, self.load, self.ts, .05, 0)

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