import calendar
from collections import defaultdict
from datetime import date

from django.db.models import Count
from django.db.models.functions import TruncMonth

from apps.integrations.models import HISPatientVisit

PATIENT_TYPE_MAP = {
    "OP": "OPD",
    "ED": "ER",
    "IP": "IP",
}

AREA_ORDER = ["OPD", "ER", "IP"]


class CensusService:
    def __init__(self, hospital_id, year=None):
        self.hospital_id = hospital_id
        self.year = year or date.today().year

    def get_monthly_counts(self, year=None):
        year = year if year is not None else self.year
        qs = (
            HISPatientVisit.objects.filter(
                hospital_id=self.hospital_id,
                admit_date__year=year,
            )
            .annotate(month=TruncMonth("admit_date"))
            .values("month", "patient_type")
            .annotate(count=Count("id"))
            .order_by("month", "patient_type")
        )

        data = {}
        for row in qs:
            m = row["month"].month
            label = PATIENT_TYPE_MAP.get(row["patient_type"], row["patient_type"])
            if m not in data:
                data[m] = {"OPD": 0, "ER": 0, "IP": 0}
            data[m][label] = row["count"]
        return data

    def get_quarterly_data(self, year=None):
        year = year if year is not None else self.year
        monthly = self.get_monthly_counts(year)
        quarters = {}
        for q in range(1, 5):
            months_in_q = [(q - 1) * 3 + i for i in range(1, 4)]
            q_data = {"months": {}, "totals": {"OPD": 0, "ER": 0, "IP": 0}}
            for m in months_in_q:
                m_data = monthly.get(m, {"OPD": 0, "ER": 0, "IP": 0})
                q_data["months"][m] = m_data
                for area in AREA_ORDER:
                    q_data["totals"][area] += m_data[area]
            quarters[q] = q_data
        return quarters

    def get_year_totals(self, year=None):
        year = year if year is not None else self.year
        monthly = self.get_monthly_counts(year)
        totals = {"OPD": 0, "ER": 0, "IP": 0, "total": 0}
        for m_data in monthly.values():
            for area in AREA_ORDER:
                totals[area] += m_data[area]
        totals["total"] = sum(totals[a] for a in AREA_ORDER)
        return totals

    def get_comparison_data(self, years=None):
        if years is None:
            years = [self.year - 2, self.year - 1, self.year]
        comparison = {}
        for y in years:
            comparison[y] = {
                "quarterly": self.get_quarterly_data(y),
                "yearly": self.get_year_totals(y),
            }
        return comparison

    def get_available_years(self):
        years = (
            HISPatientVisit.objects.filter(hospital_id=self.hospital_id)
            .values_list("admit_date__year", flat=True)
            .distinct()
            .order_by("admit_date__year")
        )
        return list(years)

    def get_chart_data(self, years=None):
        if years is None:
            years = [self.year - 2, self.year - 1, self.year]
        years = sorted(years)

        yoy_data = {area: {y: [] for y in years} for area in AREA_ORDER}
        for y in years:
            q_data = self.get_quarterly_data(y)
            for q in range(1, 5):
                for area in AREA_ORDER:
                    yoy_data[area][y].append(q_data[q]["totals"][area])

        monthly_data = {area: [] for area in AREA_ORDER}
        monthly = self.get_monthly_counts(self.year)
        for m in range(1, 13):
            m_data = monthly.get(m, {"OPD": 0, "ER": 0, "IP": 0})
            for area in AREA_ORDER:
                monthly_data[area].append(m_data[area])

        return {
            "yoy": yoy_data,
            "monthly": monthly_data,
            "quarters": ["Q1", "Q2", "Q3", "Q4"],
            "months": [calendar.month_abbr[i] for i in range(1, 13)],
            "years": years,
        }