from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.naive_bayes import MultinomialNB
import joblib
from preprocessing import preprocess_text

# ==========================================================
# 📊 DATASET LATIH REALISTIS KECAMATAN WONOSALAM (120 DATA)
# ==========================================================
raw_data = [
    # --- INFRASTRUKTUR & FASILITAS (20 Data) ---
    ("Jalan di desa Botorejo rusak parah dan berlubang besar membahayakan", "infrastruktur_fasilitas", "penting"),
    ("Jembatan penghubung desa Bunderan retak dan hampir roboh diterjang banjir", "infrastruktur_fasilitas", "penting"),
    ("Aspal di jalan raya Wonosalam mengelupas dan banyak kerikil tajam berserakan", "infrastruktur_fasilitas", "biasa"),
    ("Fasilitas gazebo di balai desa Wonosalam rusak dan tidak terawat", "infrastruktur_fasilitas", "biasa"),
    ("Tembok penahan tanah di desa Doreng longsor menutupi sebagian jalan", "infrastruktur_fasilitas", "penting"),
    ("Pagar pembatas jalan di tikungan Getas roboh akibat tertabrak truk", "infrastruktur_fasilitas", "biasa"),
    ("Jalan setapak menuju sawah di desa Jogoloyo becek dan berlumpur sulit dilewati", "infrastruktur_fasilitas", "biasa"),
    ("Gedung balai pertemuan desa Kalianyar atapnya bocor saat hujan deras", "infrastruktur_fasilitas", "biasa"),
    ("Saluran irigasi pertanian di desa Karangrejo jebol menyumbat aliran air", "infrastruktur_fasilitas", "penting"),
    ("Jalan aspal arah lapangan desa Karangrowo retak-retak berisiko longsor", "infrastruktur_fasilitas", "penting"),
    ("Fasilitas toilet umum di pasar Wonosalam rusak parah dan tidak bisa digunakan", "infrastruktur_fasilitas", "biasa"),
    ("Jembatan gantung di desa Doreng tali slingnya berkarat dan bergoyang ekstrem", "infrastruktur_fasilitas", "penting"),
    ("Akses jalan desa Botorejo berbatu-batu tajam merusak ban kendaraan warga", "infrastruktur_fasilitas", "biasa"),
    ("Tiang pembatas jalan di desa Bunderan pecah-pecah", "infrastruktur_fasilitas", "biasa"),
    ("Paving block halaman Puskesmas Wonosalam amblas membuat genangan air", "infrastruktur_fasilitas", "biasa"),
    ("Jembatan kayu di Jogoloyo rapuh dan kayunya mulai busuk membahayakan", "infrastruktur_fasilitas", "penting"),
    ("Kerusakan aspal jalan tanjakan desa Kalianyar membuat motor sering mogok", "infrastruktur_fasilitas", "biasa"),
    ("Saluran drainase air di Wonosalam ambrol menyumbat gorong-gorong", "infrastruktur_fasilitas", "penting"),
    ("Tebing jalan raya Getas longsor kecil kerikil jatuh ke jalan raya", "infrastruktur_fasilitas", "penting"),
    ("Taman bermain anak di desa Karangrowo rusak berkarat dan tajam", "infrastruktur_fasilitas", "biasa"),

    # --- KEBERSIHAN & LINGKUNGAN (20 Data) ---
    ("Tumpukan sampah liar di pinggir hutan pinus Getas bau busuk", "kebersihan_lingkungan", "biasa"),
    ("Aliran sungai di Jogoloyo tersumbat sampah plastik rawan banjir bandang", "kebersihan_lingkungan", "penting"),
    ("Warga buang sampah sembarangan di bantaran sungai desa Bunderan kotor", "kebersihan_lingkungan", "biasa"),
    ("Sampah rumah tangga menumpuk di TPS Wonosalam tidak diangkut tiga minggu", "kebersihan_lingkungan", "biasa"),
    ("Limbah cair dari home industry tahu mencemari air sungai desa Kalianyar", "kebersihan_lingkungan", "penting"),
    ("Hutan desa di Doreng dikotori botol plastik oleh wisatawan yang berkemah", "kebersihan_lingkungan", "biasa"),
    ("Bau menyengat dari bangkai hewan dibuang sembarangan di selokan Botorejo", "kebersihan_lingkungan", "biasa"),
    ("Saluran air pemukiman di Karangrejo macet tertimbun lumpur dan sampah daun", "kebersihan_lingkungan", "biasa"),
    ("Penebangan pohon liar di kawasan hutan lindung Wonosalam merusak lingkungan", "kebersihan_lingkungan", "penting"),
    ("Limbah kotoran sapi perah dialirkan langsung ke selokan desa Botorejo", "kebersihan_lingkungan", "penting"),
    ("Penumpukan sampah plastik di area balai desa Bunderan merusak pemandangan", "kebersihan_lingkungan", "biasa"),
    ("Sungai di desa Jogoloyo keruh dan banyak busa akibat limbah detergen warga", "kebersihan_lingkungan", "biasa"),
    ("Tanah longsor akibat pembalakan liar di tebing desa Doreng sangat mengkhawatirkan", "kebersihan_lingkungan", "penting"),
    ("Lumpur sisa galian tanah memenuhi jalan desa Getas licin membahayakan", "kebersihan_lingkungan", "penting"),
    ("Asap pembakaran sampah warga di dekat Puskesmas Wonosalam bikin sesak napas", "kebersihan_lingkungan", "penting"),
    ("Banyak lalat akibat sampah organik restoran menumpuk di sekitar pasar Wonosalam", "kebersihan_lingkungan", "biasa"),
    ("Selokan tersumbat sampah popok bayi di desa Karangrejo air meluap ke jalan", "kebersihan_lingkungan", "biasa"),
    ("Bau kotoran kambing menyengat dekat pemukiman padat desa Karangrowo mengganggu", "kebersihan_lingkungan", "biasa"),
    ("Wisata lapangan olahraga Kalianyar kotor berserakan sampah plastik sisa makanan", "kebersihan_lingkungan", "biasa"),
    ("Bencana longsor tebing di Wonosalam menutup aliran sungai irigasi warga", "kebersihan_lingkungan", "penting"),

    # --- PELAYANAN ADMINISTRASI (20 Data) ---
    ("Pelayanan pembuatan KTP di kantor kecamatan Wonosalam sangat lambat antrean lama", "pelayanan_administrasi", "biasa"),
    ("Pegawai kantor kelurahan Karangrejo tidak ramah dan mempersulit urus domisili", "pelayanan_administrasi", "biasa"),
    ("Petugas administrasi desa Doreng sering terlambat datang kantor kosong jam sembilan", "pelayanan_administrasi", "biasa"),
    ("Pengurusan surat waris di kecamatan bertele-tele dipersulit mintak pungli", "pelayanan_administrasi", "penting"),
    ("Sistem antrean online di kelurahan Bunderan sering error dan tidak berfungsi", "pelayanan_administrasi", "biasa"),
    ("Sulit menemui kepala desa Kalianyar untuk minta tanda tangan berkas penting", "pelayanan_administrasi", "biasa"),
    ("Proses pencairan dana desa dinilai lambat menghambat program pembangunan", "pelayanan_administrasi", "penting"),
    ("Pengurusan KK baru di kelurahan Getas memakan waktu berbulan-bulan tanpa kejelasan", "pelayanan_administrasi", "biasa"),
    ("Loket pelayanan desa Jogoloyo tutup sebelum jamnya petugas sudah pulang", "pelayanan_administrasi", "biasa"),
    ("Pegawai kelurahan Botorejo kurang paham alur pendaftaran izin usaha", "pelayanan_administrasi", "biasa"),
    ("Pengaduan warga ke kantor kecamatan Wonosalam tidak pernah ditanggapi atau direspon", "pelayanan_administrasi", "biasa"),
    ("Biaya pembuatan sertifikat tanah di desa Karangrowo tidak transparan ada pungutan liar", "pelayanan_administrasi", "penting"),
    ("Sikap diskriminatif oknum perangkat desa Doreng dalam melayani warga miskin", "pelayanan_administrasi", "penting"),
    ("Server pembuatan KTP elektronik di Wonosalam mati sudah tiga hari warga kecewa", "pelayanan_administrasi", "penting"),
    ("Syarat pengurusan kartu jaminan kesehatan di kelurahan berbelit-belit membingungkan", "pelayanan_administrasi", "biasa"),
    ("Informasi pelayanan di kantor desa Bunderan tidak jelas tidak ada papan pengumuman", "pelayanan_administrasi", "biasa"),
    ("Petugas kelurahan Kalianyar salah mengetik nama di akta kelahiran anak saya", "pelayanan_administrasi", "biasa"),
    ("Urus surat izin keramaian di polsek Wonosalam lambat dan birokrasinya ribet", "pelayanan_administrasi", "biasa"),
    ("Waktu tunggu legalisir berkas di kecamatan Wonosalam tidak pasti sangat lama", "pelayanan_administrasi", "biasa"),
    ("Oknum perangkat desa Botorejo mempersulit warga mendapatkan surat pengantar nikah", "pelayanan_administrasi", "penting"),

    # --- LAYANAN PUBLIK & UTILITAS (20 Data) ---
    ("Lampu penerangan jalan umum di desa Doreng padam sudah lebih dari seminggu", "layanan_publik_utilitas", "penting"),
    ("Pasokan air bersih dari HIPPAM di desa Kalianyar mampet warga sulit mandi", "layanan_publik_utilitas", "penting"),
    ("Sinyal internet seluler di desa Botorejo sangat lemah mengganggu sekolah daring", "layanan_publik_utilitas", "biasa"),
    ("Listrik di wilayah kecamatan Wonosalam sering padam tanpa pemberitahuan dahulu", "layanan_publik_utilitas", "penting"),
    ("Air keran dari PDAM Wonosalam keruh kecoklatan berlumpur tidak layak konsumsi", "layanan_publik_utilitas", "penting"),
    ("Kabel listrik PLN menjuntai hampir menyentuh tanah di Bunderan bahaya", "layanan_publik_utilitas", "penting"),
    ("Jaringan telepon seluler mati total setelah hujan badai di Getas", "layanan_publik_utilitas", "biasa"),
    ("Pipa air bersih bocor di pinggir jalan desa Jogoloyo air terbuang percuma", "layanan_publik_utilitas", "biasa"),
    ("Penerangan jalan raya arah Wonosalam minim rawan kecelakaan di malam hari", "layanan_publik_utilitas", "penting"),
    ("Kabel internet wifi bergelantungan berantakan di pohon jalan desa Karangrowo", "layanan_publik_utilitas", "biasa"),
    ("Pompa air sumur bor desa Karangrejo rusak membuat warga kekeringan air bersih", "layanan_publik_utilitas", "penting"),
    ("Trafo listrik meledak di desa Doreng mengakibatkan pemadaman listrik massal", "layanan_publik_utilitas", "penting"),
    ("Saluran air bersih warga tersumbat pasir di desa Botorejo HIPPAM mati", "layanan_publik_utilitas", "biasa"),
    ("Sinyal handphone hilang saat mendung tebal di pelosok Kalianyar terisolasi", "layanan_publik_utilitas", "biasa"),
    ("Kabel listrik PLN tertimpa dahan pohon tumbang di desa Wonosalam percikan api", "layanan_publik_utilitas", "penting"),
    ("Penyaluran air bersih desa Bunderan digilir seminggu hanya mengalir dua kali", "layanan_publik_utilitas", "biasa"),
    ("Lampu JPU mati di area tikungan tajam Getas sangat gelap menakutkan", "layanan_publik_utilitas", "penting"),
    ("Tagihan air HIPPAM desa Jogoloyo melonjak tidak wajar padahal air jarang mengalir", "layanan_publik_utilitas", "biasa"),
    ("Tiang listrik PLN di pinggir jalan desa Karangrowo miring rawan roboh ke jalan", "layanan_publik_utilitas", "penting"),
    ("Gangguan jaringan kabel telepon di kantor kecamatan Wonosalam menyulitkan komunikasi", "layanan_publik_utilitas", "biasa"),

    # --- KEAMANAN & KETERTIBAN (20 Data) ---
    ("Terjadi aksi pencurian sepeda motor di area parkir Wonosalam tadi malam", "keamanan_ketertiban", "penting"),
    ("Banyak pemuda mabuk-mabukan di pinggir jalan desa Jogoloyo bikin resah warga", "keamanan_ketertiban", "penting"),
    ("Aksi balap liar sepeda motor knalpot brong di jalan raya Wonosalam mengganggu", "keamanan_ketertiban", "penting"),
    ("Pencurian pisang dan hasil kebun marak terjadi di perkebunan desa Doreng", "keamanan_ketertiban", "biasa"),
    ("Pos kamling di desa Kalianyar kosong tidak pernah ada ronda malam rawan maling", "keamanan_ketertiban", "biasa"),
    ("Premanisme berkedok uang parkir liar di tempat balai desa Bunderan", "keamanan_ketertiban", "biasa"),
    ("Terjadi perkelahian antar kelompok pemuda di acara hiburan malam desa Botorejo", "keamanan_ketertiban", "penting"),
    ("Aksi penipuan hipnotis keliling menyasar ibu rumah tangga di kelurahan Karangrejo", "keamanan_ketertiban", "penting"),
    ("Pencurian ternak kambing milik warga desa Karangrowo saat dini hari bobol kandang", "keamanan_ketertiban", "penting"),
    ("Banyak anjing liar berkeliaran di Getas mengejar anak-anak sekolah takut digigit", "keamanan_ketertiban", "biasa"),
    ("Peredaran obat-obatan terlarang pil koplo dicurigai menyasar remaja Wonosalam", "keamanan_ketertiban", "penting"),
    ("Aksi corat-coret vandalisme merusak fasilitas umum di pusat kecamatan Wonosalam", "keamanan_ketertiban", "biasa"),
    ("Suara petasan keras saat malam hari mengganggu istirahat warga desa Jogoloyo", "keamanan_ketertiban", "biasa"),
    ("Pencurian kotak amal di masjid desa Doreng terekam kamera CCTV pelaku kabur", "keamanan_ketertiban", "penting"),
    ("Keributan antar tetangga akibat sengketa batas tanah di desa Bunderan memanas", "keamanan_ketertiban", "biasa"),
    ("Warga mencurigai rumah kosong dijadikan tempat transaksi barang curian di Kalianyar", "keamanan_ketertiban", "penting"),
    ("Aksi pemalakan sopir angkutan barang di perbatasan Wonosalam-Kandangan meresahkan", "keamanan_ketertiban", "penting"),
    ("Banyak kendaraan parkir sembarangan di jalan pasar Wonosalam memicu kemacetan parah", "keamanan_ketertiban", "biasa"),
    ("Pembobolan warung kelontong di malam hari kehilangan rokok dan sembako di Getas", "keamanan_ketertiban", "penting"),
    ("Aktivitas penambangan liar batu sungai di Botorejo merusak pondasi jembatan", "keamanan_ketertiban", "penting"),

    # --- SOSIAL & KESEJAHTERAAN (20 Data) ---
    ("Bantuan sosial PKH di desa Karangrowo dinilai tidak tepat sasaran diberikan orang kaya", "sosial_kesejahteraan", "biasa"),
    ("Warga lansia miskin sebatang kara di desa Doreng butuh bantuan sembako dan obat", "sosial_kesejahteraan", "penting"),
    ("Penyaluran BLT Dana Desa di kelurahan Kalianyar ricuh antrean tidak tertib", "sosial_kesejahteraan", "biasa"),
    ("Banyak anak putus sekolah di desa Botorejo karena orang tua tidak mampu biaya", "sosial_kesejahteraan", "penting"),
    ("Warga kurang mampu menderita sakit parah di Bunderan tidak punya biaya berobat", "sosial_kesejahteraan", "penting"),
    ("Distribusi kartu Indonesia sehat KIS tidak merata banyak warga miskin terlewat", "sosial_kesejahteraan", "biasa"),
    ("Gelandangan dan pengemis mulai marak di sekitar pertigaan pusat Wonosalam", "sosial_kesejahteraan", "biasa"),
    ("Program bedah rumah warga miskin di desa Jogoloyo mangkrak pengerjaan terhenti", "sosial_kesejahteraan", "biasa"),
    ("Warga penyandang disabilitas di desa Karangrejo membutuhkan bantuan kursi roda", "sosial_kesejahteraan", "biasa"),
    ("Kelaparan warga akibat gagal panen cengkeh diserang hama di lereng Wonosalam", "sosial_kesejahteraan", "penting"),
    ("Kenaikan harga beras di pasar Wonosalam mencekik warga ekonomi menengah ke bawah", "sosial_kesejahteraan", "biasa"),
    ("Anak yatim piatu di desa Getas butuh uluran tangan untuk biaya sekolah dasar", "sosial_kesejahteraan", "biasa"),
    ("Penerima bantuan pangan nontunai BPNT desa Karangrowo dicurangi berasnya kualitas jelek bau", "sosial_kesejahteraan", "penting"),
    ("Warga miskin tidak terdata di DTKS sehingga tidak pernah menerima bantuan apapun", "sosial_kesejahteraan", "biasa"),
    ("Kasus gizi buruk stunting ditemukan pada anak balita keluarga miskin Botorejo", "sosial_kesejahteraan", "penting"),
    ("Bencana kebakaran rumah melanda warga desa Doreng kehilangan tempat tinggal butuh tenda", "sosial_kesejahteraan", "penting"),
    ("Pekerja buruh kebun cengkeh Wonosalam di-PHK massal menganggur tidak ada pendapatan", "sosial_kesejahteraan", "biasa"),
    ("Warga jompo terlantar hidup di gubuk reyot desa Kalianyar butuh perhatian dinsos", "sosial_kesejahteraan", "penting"),
    ("Bantuan alat pertanian handtractor di desa Jogoloyo dikuasai pribadi oleh kelompok tani", "sosial_kesejahteraan", "biasa"),
    ("Penderita penyakit menular TBC di desa Karangrejo butuh isolasi dan bantuan pangan mendesak", "sosial_kesejahteraan", "penting")
]

# ==========================================================
# ⚙️ RUN TRAINING (ONLY WHEN RUN DIRECTLY)
# ==========================================================
if __name__ == '__main__':
    from preprocessed_dataset import preprocessed_data
    print("Memuat Dataset Prekomputasi...")
    texts = []
    categories = []
    priorities = []

    for text, clean_text, cat, prio in preprocessed_data:
        texts.append(clean_text)
        categories.append(cat)
        priorities.append(prio)

    print("Pemuatan dataset selesai!")

    # ==========================================================
    # 🔮 TRAINING TF-IDF + MULTINOMIAL NAIVE BAYES
    # ==========================================================
    # TF-IDF Vectorizer
    print("Training TF-IDF Vectorizer...")
    tfidf = TfidfVectorizer(
        lowercase=True,
        ngram_range=(1,1)
    )
    X = tfidf.fit_transform(texts)

    # Model Kategori
    print("Training Model Naive Bayes - Kategori...")
    model_kategori = MultinomialNB(alpha=0.01)
    model_kategori.fit(X, categories)

    # Model Prioritas
    print("Training Model Naive Bayes - Prioritas...")
    model_prioritas = MultinomialNB(alpha=0.01)
    model_prioritas.fit(X, priorities)

    # ==========================================================
    # 💾 SIMPAN MODEL (.PKL)
    # ==========================================================
    print("Menyimpan model ke disk...")
    joblib.dump(model_kategori, 'model_kategori.pkl')
    joblib.dump(model_prioritas, 'model_prioritas.pkl')
    joblib.dump(tfidf, 'tfidf.pkl')

    print("Model berhasil diperbarui dan disimpan dengan 120 dataset Wonosalam!")