import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns

plt.rcParams = plt.rcParamsOrig

"""
Explanation: Artificial Intelligence & Machine Learning
Tugas 4: Bayesian Networks & Natural Language Processing
Mekanisme
Anda hanya diwajibkan untuk mengumpulkan file ini saja ke uploader yang disediakan di https://elearning.uai.ac.id/. Ganti nama file ini saat pengumpulan menjadi tugas4_NIM.ipynb.
Keterlambatan: Pengumpulan tugas yang melebihi tenggat yang telah ditentukan tidak akan diterima. Keterlambatan akan berakibat pada nilai nol untuk tugas ini.
Kolaborasi: Anda diperbolehkan untuk berdiskusi dengan teman Anda, tetapi dilarang keras menyalin kode maupun tulisan dari teman Anda. Kecurangan akan berakibat pada nilai nol untuk tugas ini.
Petunjuk
Anda diperbolehkan (jika dirasa perlu) untuk mengimpor modul tambahan untuk tugas ini. Namun, seharusnya modul yang tersedia sudah cukup untuk memenuhi kebutuhan Anda. Untuk kode yang Anda ambil dari sumber lain, cantumkan URL menuju referensi tersebut jika diambil dari internet!
Perhatikan poin untuk tiap soal! Semakin kecil poinnya, berarti kode yang diperlukan untuk menjawab soal tersebut seharusnya semakin sedikit!
End of explanation
"""


from tqdm.notebook import trange

class HMM:
    def __init__(self, p_start, p_trans, p_emit, p_stop=None):
        assert p_trans.shape[0] == p_emit.shape[0]
        self.p_start = p_start
        self.p_trans = p_trans
        self.p_emit = p_emit
        if p_stop is not None:
            self.p_stop = p_stop
        else:
            self.p_stop = np.ones(self.p_trans.shape[0]) # mengisi dengan 1 untuk probabilitas berhenti
    
    def forward(self, sequence: list):
        a = []
        for i, x in enumerate(sequence):
            if i == 0:
                a.append((self.p_start * self.p_emit[x]).values)
            else:
                a.append(
                    a[i-1].dot(
                        self.p_trans.dot(np.diag(self.p_emit[x]))
                    )
                )
        self.a = np.array(a)
        return self
    
    def backward(self, sequence: list):
        b = []
        prev_x = sequence[-1]
        for i, x in enumerate(sequence[::-1]): # membaca observasi dari belakang ke depan
            if i == 0:
                b.append(self.p_stop)
            else:
                b.append(
                    b[i-1].dot(
                        np.diag(self.p_emit[prev_x]).dot(
                            self.p_trans.T
                        )
                    )
                )
            prev_x = x
        self.b = np.array(b[::-1]) # mengembalikan urutan β
        return self
    
    def forward_backward(self, sequence: list):
        return self.forward(sequence).backward(sequence).normalised_gamma
    
    @property
    def gamma(self):
        """γ = α * β
        Mengembalikan nilai γ di setiap state.
        """
        pass
    
    @property
    def normalised_gamma(self):
        """Normalisasi γ agar mendapatkan probabilitas.
        Total semua nilai dalam satu baris = 1.
        """
        pass
    
    def log_probability(self, seq: list) -> float:
        """Menjalankan forward-backward agar mendapatkan nilai γ.
        Mengembalikan nilai log dari jumlah nilai γ dari salah satu state.
        
        Catatan: Total nilai γ di setiap state selalu sama.
        """
        pass

    def fit(self, seq: list, max_iter=100):
        """Memperbaiki nilai p_emit berdasarkan observasi yang dilakukan.
        Tidak mengembalikan nilai apapun.
        """
        pass

    def predict(self, seq):
        return self.forward_backward(seq).argmax(axis=1)
    
    def predict_proba(self, seq):
        return self.forward_backward(seq)

"""
Explanation: 1. Hidden Markov Model (10 poin)
Soal 1.1 (8 poin)
Anda diberikan kelas Hidden Markov Model yang harus Anda lengkapi bagian-bagian di dalamnya. Dalam mengisi bagian-bagian tersebut, Anda tidak diperkenankan mengubah kerangka yang sudah diberikan.
Implementasikan metode/fungsi:
1. gamma (2 poin)
2. normalised_gamma (2 poin)
3. log_probability (2 poin)
4. fit (2 poin)
Untuk memudahkan, dalam implementasi di bawah ini variabel a adalah α dan variabel b adalah β.
End of explanation
"""


# Kasus uji - jangan diubah!
p_emit = pd.DataFrame({
    's1': {'A': 0.25, 'C': 0.25, 'G': 0.25, 'T': 0.25},
    's2': {'A': 0.10, 'C': 0.40, 'G': 0.40, 'T': 0.10}
}).T
p_trans = pd.DataFrame({
    's1': {'s1': 0.89, 's2': 0.10},
    's2': {'s1': 0.1, 's2': 0.9}
}).T
mdl = HMM(
    p_start=np.array([0.5, 0.5]),
    p_trans=p_trans,
    p_emit=p_emit,
    p_stop=np.array([0.01, 0.])
)

seq = list('CGACTACTGACTACTCGCCGACGCGACTGCCGTCTATACTGCGCATACGGC')
hmm_predictions = mdl.predict(seq)
print("sequence: {}".format(''.join(seq)))
print("hmm pred: {}".format(''.join(map(str, hmm_predictions))))

assert ''.join(map(str, hmm_predictions)) == "000000000000000111111111111111100000000000000000000"
assert np.allclose(mdl.forward_backward(seq).sum(axis=1), np.ones(len(seq)))
assert np.isclose(mdl.log_probability(seq), -76.1486)

"""
Explanation: Contoh kasus di bawah ini diadaptasi dari sini. Contoh ini adalah penggunaan HMM untuk mencari sequence DNA yang banyak mengandung nukleotida CG. Contoh kasus ini merupakan penyederhanaan dari kasus biologi yang riil dalam DNA sequencing.
Anda dapat menggunakan kode di bawah ini untuk menguji hasil implementasi Anda. Kalau tidak ada pesan error, berarti implementasi Anda sudah benar.
End of explanation
"""


import string

def translate(states: list):
    pass

"""
Explanation: Soal 1.2 (2 poin)
Dari contoh kasus uji di atas, jelaskan maksud nilai p_start = (0.5, 0.5)!
Jawaban Anda di sini
Soal 1.3.a (bonus - 2 poin)
Lengkapi fungsi berikut untuk menerjemahkan state dari HMM ke karakter alfabet dan spasi.
End of explanation
"""


cipher_text = 'cipx qidnsi xdyedlsxi yldnid xdx pldtizicid clplwylciid xdyedlsxi bim bim tidn pldnldix vlpxdyibid clcuisiid yid mixd mixd yxslmldnniwicid yldnid aiwi slcsipi yid yimip zlpve tidn slsxdnciz sxdncizdti'
K = 27
p_trans = pd.read_csv('https://raw.githubusercontent.com/aliakbars/uai-ai/master/datasets/p_trans.csv', index_col=0)

# Kode Anda di sini

"""
Explanation: Soal 1.3.b (bonus - 3 poin)
Dengan menggunakan HMM, terjemahkan cipher text yang diberikan berikut ini. Probabilitas transisi (p_trans) yang dilatih dari 10000 teks berita berbahasa Indonesia sudah diberikan untuk Anda. Sebagai petunjuk, Anda dapat menggunakan uniform distribution untuk p_start dan p_emit di awal.
Catatan: Hasil prediksi dari model HMM tidak akan 100% benar. 1 poin dari soal ini didapat dari perbaikan prediksi HMM ke teks asli (plain text) dari cipher text tersebut secara manual.
End of explanation
"""


import re

articles = pd.read_csv(
    'https://raw.githubusercontent.com/aliakbars/uai-ai/master/datasets/ind_news_2012_10K-sentences.txt',
    sep='\t', header=None, quoting=3
)[1][:2000]
articles = articles.str.lower().apply(lambda x: re.findall("[a-z]+", x))
articles

"""
Explanation: 2. Natural Language Processing (10 poin)
Natural language processing (NLP) berkaitan cukup erat dengan Hidden Markov Model (HMM) dan Bayesian networks secara umum. Salah satu hubungannya adalah konsep language modelling. Di soal di atas, Anda sudah diberikan langsung p_trans untuk digunakan. Pada bagian ini, Anda diminta untuk melatih probabilistic language models, i.e. p_trans, untuk "mengajari" komputer/mesin untuk "mengerti" bahasa manusia. Artikel yang digunakan dalam tugas ini diambil dari Leipzig Corpora Collection.
End of explanation
"""


from collections import Counter
from nltk import ngrams

# Kode Anda di sini

"""
Explanation: Soal 2.1 (2 poin)
Salah satu konsep dasar dalam language model adalah n-gram. Secara matematis, persoalan ini diformulasikan sebagai:
$$
p(w_1,w_2,...,w_{t-1},w_t) \approx \prod_{i=1}^{t} p(w_i|w_{i-1})
$$
sehingga kita bisa memanfaatkan bigram (2-gram) untuk memperdiksi kata berikutnya jika diketahui kata sekarang. Hasil pencacahan dari kumpulan bigram itulah yang dapat dijadikan sebagai probabilitas $p(w_t|w_{t-1})$.
$$
p(w_2|w_1) = \frac{c(w_1, w_2)}{c(w_1)}
$$
Petunjuk: Referensi tambahan bisa dilihat di buku ini hal. 4. Anda boleh memanfaatkan modul Counter dan ngrams. Namun, looping bisa menjadi solusi yang lebih sederhana untuk Anda.
End of explanation
"""


# Kode Anda di sini

"""
Explanation: Soal 2.2 (2 poin)
Untuk menyederhanakan, ambil 5000 bigrams yang paling sering muncul. Lalu, buatlah pandas DataFrame dalam variabel df yang berbentuk:
|    | w1      | w2      |   freq |
|---:|:--------|:--------|-------:|
|  0 | antara  | news    |     62 |
|  1 | saat    | ini     |     52 |
|  2 | salah   | satu    |     42 |
|  3 | menurut | dia     |     41 |
|  4 | di      | jakarta |     34 |
End of explanation
"""


# Kode ini digunakan untuk Laplace smoothing
dfc = (
    df
    .set_index(['w1','w2'])
    .freq
    .unstack(fill_value=0)
    .stack()
    .reset_index(name='freq')
)
dfc['freq'] += 1

def get_next_word(dfc, w):
    pass

"""
Explanation: Soal 2.3 (2 poin)
Lengkapi fungsi berikut yang dapat memilih kata $w_t$ yang probabilitasnya paling besar jika diberikan $w_{t-1}$, i.e.
$$
w_t = \arg \max_{w_i} p(w_i|w_{t-1})
$$
End of explanation
"""


def get_next_random_word(dfc, w, n=5):
    pass

"""
Explanation: Soal 2.4 (2 poin)
Sebagai variasi, Anda dapat mengganti kata yang dipilih untuk $w_t$ dengan memilih secara acak dari 5 kata dengan probabilitas tertinggi.
End of explanation
"""


def generate_sentence(dfc, n=10):
    pass

generate_sentence(dfc)

"""
Explanation: Soal 2.5 (2 poin)
Buatlah fungsi yang menerima input berupa dfc dan jumlah kata dalam kalimat yang mau dihasilkan, serta menggunakan get_next_word() untuk menghasilkan kata berikutnya.
End of explanation
"""


# Kode Anda di sini

"""
Explanation: Soal 2.6 (bonus - 5 poin)
Dari fungsi generate_sentence() di atas, kata pertama selalu dipilih secara acak dan berakhir setelah jumlah kata sesuai parameter n. Modifikasi representasi bigram agar dapat menerima &lt;start&gt; dan &lt;end&gt; sebagai token untuk mengawali dan mengakhiri kalimat. Dengan demikian, panjang kalimat akan ditentukan oleh kata terakhir yang muncul dan awal kalimat dimulai dengan:
1. kata-kata yang paling sering dipakai untuk mengawali kalimat di sebuah artikel; atau
2. diberikan oleh pengguna.
End of explanation
"""


