from experta import KnowledgeEngine, Rule, Fact, AS

ID_TO_NAME = {
    "1": "ho_kho", "2": "ho_mau", "3": "sot_nhe", "4": "mo_hoi", "5": "sut_can",
    "6": "met_moi", "7": "tiep_xuc_lao", "8": "hut_thuoc", "9": "tiep_xuc_khoi_bui",
    "10": "ho_dam_sang", "11": "kho_tho", "12": "viem_phe_quan", "13": "tren_40_tuoi"
}

class TBKnowledge(KnowledgeEngine):
    # === LAO PHỔI ===
    @Rule(Fact(ho_kho=True), Fact(sot_nhe=True), Fact(mo_hoi=True))
    def nghi_ngo_lao_tam_chung(self): self.declare(Fact(nghi_ngo="lao"))

    @Rule(Fact(ho_mau=True))
    def nghi_ngo_lao_ho_mau(self): self.declare(Fact(nghi_ngo="lao"))

    @Rule(Fact(sut_can=True), Fact(met_moi=True))
    def nghi_ngo_lao_suy_sut(self): self.declare(Fact(nghi_ngo="lao"))

    @Rule(Fact(tiep_xuc_lao=True))
    def nghi_ngo_lao_tiep_xuc(self): self.declare(Fact(nghi_ngo="lao"))

    @Rule(Fact(ho_kho=True), Fact(sot_nhe=True))
    def nghi_ngo_lao_ket_hop(self): self.declare(Fact(nghi_ngo="lao"))

    # === COPD ===
    @Rule(Fact(hut_thuoc=True), Fact(ho_dam_sang=True))
    def nghi_ngo_copd_hut_ho(self): self.declare(Fact(nghi_ngo="copd"))

    @Rule(Fact(kho_tho=True), Fact(viem_phe_quan=True))
    def nghi_ngo_copd_kho_viem(self): self.declare(Fact(nghi_ngo="copd"))

    @Rule(Fact(tiep_xuc_khoi_bui=True), Fact(tren_40_tuoi=True))
    def nghi_ngo_copd_nghe_nghiep(self): self.declare(Fact(nghi_ngo="copd"))

    @Rule(Fact(hut_thuoc=True), Fact(kho_tho=True), Fact(tren_40_tuoi=True))
    def nghi_ngo_copd_nang(self): self.declare(Fact(nghi_ngo="copd"))

    @Rule(Fact(ho_dam_sang=True), Fact(tren_40_tuoi=True))
    def nghi_ngo_copd_ho_lau(self): self.declare(Fact(nghi_ngo="copd"))

    # === GHÉP CÂU NÂNG CAO ===
    @Rule(Fact(ho_kho=True), Fact(mo_hoi=True), Fact(tiep_xuc_lao=True))
    def nghi_ngo_lao_ca_nang(self): self.declare(Fact(nghi_ngo="lao"))

def run_rule_engine(answers):
    engine = TBKnowledge()
    engine.reset()
    for a in answers:
        if a.answer and str(a.questionnaire_id) in ID_TO_NAME:
            name = ID_TO_NAME[str(a.questionnaire_id)]
            engine.declare(Fact(**{name: True}))
    engine.run()
    nghi_ngo = [f["nghi_ngo"] for f in engine.facts.values() if isinstance(f, Fact) and "nghi_ngo" in f]
    return list(set(nghi_ngo))