import numpy as np
from hmmlearn import hmm

 # Number of hidden states (e.g., 0: Sunny, 1: Rainy)
n_states = 2
# 1. Define Training Data (sequence of observations)
train_obs = np.array([[0], [1], [2], [0], [1], [1], [2], [0], [0], [1]])

# 2. 
model = hmm.CategoricalHMM(n_components=n_states)

# 3. Fit 
model.fit(train_obs)
print("Learned Transition Matrix:\n", np.round(model.transmat_, 3))
print("\nLearned Emission Matrix:\n", np.round(model.emissionprob_, 3))

n_states = 2 
n_observations = 3 

start_prob = np.array([0.6, 0.4])

# P(next_state | current_state)
trans_prob = np.array([
  [0.7, 0.3], # Sunny -> (Sunny, Rainy)
  [0.4, 0.6]  # Rainy -> (Sunny, Rainy)
])

# P(observation | state)
emit_prob = np.array([
  [0.6, 0.3, 0.1], # Sunny: (Walk, Shop, Clean)
  [0.1, 0.4, 0.5]  # Rainy: (Walk, Shop, Clean)
])
model = hmm.CategoricalHMM(n_components=n_states)

model.startprob_ = start_prob
model.transmat_ = trans_prob
model.emissionprob_ = emit_prob
# Example: Walk, Shop, Clean, Walk -> [0, 1, 2, 0]
obs_seq = np.array([[0], [1], [2]]) 
# Must be 2D: (n_samples, n_features=1)
logprob, hidden_states = model.decode(obs_seq, algorithm="viterbi")
logprob = model.score(obs_seq)
print(f"Log probability of the sequence: {logprob}")

# If you still need the hidden states from Viterbi separately:
# _, hidden_states = model.decode(obs_seq, algorithm="viterbi")

# Calculate posterior probabilities P(S_t | O_1:T) for each state S_t at each time t
# This internally computes both forward and backward probabilities.
posterior_probs = model.predict_proba(obs_seq)

posterior_probs
print("--- Using Defined Model ---")
print("Observation sequence (numeric):", obs_seq.flatten())
print("Most likely hidden states (numeric):", hidden_states) # 0: Sunny, 1: Rainy
print(f"Log Probability of sequence: {logprob:.4f}")
