from datetime import date

from django.test import TestCase

from core.models import Subject, Transaction
from core.prediction import predict_transactions


class PredictionTestCase(TestCase):

    def test_not_enough_data(self):
        subject = Subject.objects.create(name="rent")
        predicted_transactions, prediction_info = predict_transactions(subject)
        self.assertEqual(predicted_transactions, [])
        self.assertIsNone(prediction_info)
        Transaction.objects.create(amount=-333, booking_date=date(2020, 6, 2), subject=subject)
        predicted_transactions, prediction_info = predict_transactions(subject)
        self.assertEqual(predicted_transactions, [])
        self.assertIsNone(prediction_info)
        Transaction.objects.create(amount=-333, booking_date=date(2020, 7, 2), subject=subject)
        predicted_transactions, prediction_info = predict_transactions(subject)
        self.assertNotEqual(predicted_transactions, [])
        self.assertIsNotNone(prediction_info)

    def test_every_month(self):
        subject = Subject.objects.create(name="rent")
        Transaction.objects.create(amount=-333, booking_date=date(2020, 6, 2), subject=subject)
        Transaction.objects.create(amount=-333, booking_date=date(2020, 7, 2), subject=subject)
        predicted_transactions, prediction_info = predict_transactions(subject)
        first, second, last = predicted_transactions[0], predicted_transactions[1], predicted_transactions[-1]
        self.assertEqual(first.booking_date, date(2020, 8, 2))
        self.assertEqual(second.booking_date, date(2020, 9, 2))
        self.assertEqual(last.booking_date, date(2021, 5, 2))
        for tr in predicted_transactions:
            self.assertEqual(tr.amount, -333)
        self.assertDictEqual(
            prediction_info,
            {"recurring_days": None, "recurring_months": 1, "day_of_month": 2},
        )

    def test_every_3_months(self):
        subject = Subject.objects.create(name="rent")
        Transaction.objects.create(amount=-333, booking_date=date(2020, 6, 2), subject=subject)
        Transaction.objects.create(amount=-333, booking_date=date(2020, 9, 2), subject=subject)
        predicted_transactions, prediction_info = predict_transactions(subject)
        first, second, last = predicted_transactions[0], predicted_transactions[1], predicted_transactions[-1]
        self.assertEqual(first.booking_date, date(2020, 12, 2))
        self.assertEqual(second.booking_date, date(2021, 3, 2))
        self.assertEqual(last.booking_date, date(2021, 9, 2))
        for tr in predicted_transactions:
            self.assertEqual(tr.amount, -333)
        self.assertDictEqual(
            prediction_info,
            {"recurring_days": None, "recurring_months": 3, "day_of_month": 2},
        )

    def test_every_year(self):
        subject = Subject.objects.create(name="rent")
        Transaction.objects.create(amount=-333, booking_date=date(2019, 6, 2), subject=subject)
        Transaction.objects.create(amount=-333, booking_date=date(2020, 6, 2), subject=subject)
        predicted_transactions, prediction_info = predict_transactions(subject)
        self.assertEqual(len(predicted_transactions), 1)
        trans = predicted_transactions[0]
        self.assertEqual(trans.booking_date, date(2021, 6, 2))
        self.assertEqual(trans.amount, -333)
        self.assertDictEqual(
            prediction_info,
            {"recurring_days": None, "recurring_months": 12, "day_of_month": 2},
        )

    def test_monthly_varying_begin_end(self):
        subject = Subject.objects.create(name="rent")
        Transaction.objects.create(amount=-333, booking_date=date(2020, 1, 31), subject=subject)
        Transaction.objects.create(amount=-333, booking_date=date(2020, 3, 1), subject=subject)
        Transaction.objects.create(amount=-333, booking_date=date(2020, 4, 1), subject=subject)
        Transaction.objects.create(amount=-333, booking_date=date(2020, 4, 30), subject=subject)
        predicted_transactions, prediction_info = predict_transactions(subject)
        first, second, last = predicted_transactions[0], predicted_transactions[1], predicted_transactions[-1]
        self.assertEqual(first.booking_date, date(2020, 5, 30))
        self.assertEqual(second.booking_date, date(2020, 6, 30))
        self.assertEqual(last.booking_date, date(2021, 2, 28))
        for tr in predicted_transactions:
            self.assertEqual(tr.amount, -333)
        self.assertDictEqual(
            prediction_info,
            {"recurring_days": None, "recurring_months": 1, "day_of_month": None},
        )

    def test_no_monthly_pattern(self):
        subject = Subject.objects.create(name="rent")
        Transaction.objects.create(amount=-333, booking_date=date(2020, 1, 15), subject=subject)
        Transaction.objects.create(amount=-333, booking_date=date(2020, 2, 10), subject=subject)
        Transaction.objects.create(amount=-333, booking_date=date(2020, 4, 25), subject=subject)
        Transaction.objects.create(amount=-333, booking_date=date(2020, 5, 5), subject=subject)
        predicted_transactions, prediction_info = predict_transactions(subject)
        first, second, last = predicted_transactions[0], predicted_transactions[1], predicted_transactions[-1]
        self.assertEqual(first.booking_date, date(2020, 6, 11))
        self.assertEqual(second.booking_date, date(2020, 7, 18))
        self.assertEqual(last.booking_date, date(2021, 5, 10))
        for tr in predicted_transactions:
            self.assertEqual(tr.amount, -333)
        self.assertDictEqual(
            prediction_info,
            {"recurring_days": 37, "recurring_months": None, "day_of_month": None},
        )

    def test_amount_change(self):
        subject = Subject.objects.create(name="rent")
        Transaction.objects.create(amount=1337, booking_date=date(2020, 6, 2), subject=subject)
        Transaction.objects.create(amount=1312, booking_date=date(2020, 7, 2), subject=subject)
        Transaction.objects.create(amount=404, booking_date=date(2020, 8, 2), subject=subject)
        predicted_transactions, prediction_info = predict_transactions(subject)
        for tr in predicted_transactions:
            self.assertEqual(tr.amount, 404)