import pgmpy
import numpy as np
from pgmpy.models import DynamicBayesianNetwork as DBN
from pgmpy.factors.discrete import TabularCPD
from pgmpy.inference import DBNInference
import networkx as nx
from flask import Flask, request, jsonify

app = Flask(__name__)

class StudentModelDBN:
    def __init__(self, skills, prerequisites, params, initial_target_skill):
        self.skills = skills
        self.prerequisites = prerequisites
        self.params = params
        self.target_skill = initial_target_skill
        self.model = DBN()
        
        # Khởi tạo belief_state ban đầu
        self.belief_state = {}
        for skill in self.skills:
            if not self.prerequisites.get(skill):
                self.belief_state[skill] = 0.5
        
        for skill in self.skills:
            if self.prerequisites.get(skill):
                prereqs = self.prerequisites[skill]
                prereq_probs = [self.belief_state[p] for p in prereqs]
                cardinality = [2] * len(prereqs)
                prob_mastered = 0.0
                for states in np.ndindex(tuple(cardinality)):
                    p_state = 1.0
                    for i, s in enumerate(states):
                        p_state *= prereq_probs[i] if s == 1 else (1 - prereq_probs[i])
                    if all(s == 1 for s in states):
                        p_m = 0.8
                    else:
                        p_m = 0.2
                    prob_mastered += p_state * p_m
                self.belief_state[skill] = prob_mastered

        # Xây dựng cấu trúc
        self._build_model_structure(self.target_skill)
        
        # Thêm xác suất (CPDs)
        self._add_conditional_probabilities(self.target_skill)
        
        # Khởi tạo trạng thái ban đầu
        self.model.initialize_initial_state()
        self.model.check_model()
        self.inference_engine = DBNInference(self.model)

    def _compute_prior(self, skill):
        """Tính prior cho một kỹ năng dựa trên prerequisites hiện có trong belief_state.
        Nếu thiếu prerequisite trong belief_state, dùng 0.5 mặc định."""
        prereqs = self.prerequisites.get(skill, [])
        if not prereqs:
            return 0.5
        prereq_probs = [self.belief_state.get(p, 0.5) for p in prereqs]
        cardinality = [2] * len(prereqs)
        prob_mastered = 0.0
        for states in np.ndindex(tuple(cardinality)):
            p_state = 1.0
            for i, s in enumerate(states):
                p_state *= prereq_probs[i] if s == 1 else (1 - prereq_probs[i])
            if all(s == 1 for s in states):
                p_m = 0.8
            else:
                p_m = 0.2
            prob_mastered += p_state * p_m
        return prob_mastered

    def _merge_external_beliefs(self, external_belief_state):
        """Hòa trộn belief_state bên ngoài (từ DB) vào mô hình hiện tại.
        Điền các kỹ năng còn thiếu bằng prior tính toán."""
        for skill in self.skills:
            if skill in external_belief_state:
                self.belief_state[skill] = external_belief_state[skill]
            elif skill not in self.belief_state:
                self.belief_state[skill] = self._compute_prior(skill)
        # Đảm bảo không còn thiếu kỹ năng nào
        missing = [s for s in self.skills if s not in self.belief_state]
        for m in missing:
            self.belief_state[m] = 0.5

    def _build_model_structure(self, target_skill):
        base_nodes = list(self.skills) + ['Correctness', 'ResponseTime']
        self.model.add_nodes_from(base_nodes)

        intra_slice_edges = []
        # Thêm edges cho skills có prerequisites
        for skill, prereqs in self.prerequisites.items():
            for prereq in prereqs:
                intra_slice_edges.append(((prereq, 0), (skill, 0)))

        # Inter-slice edges: ONLY for skills that have intra-slice connections OR are isolated
        # Lấy tất cả skills xuất hiện trong intra_slice_edges
        connected_skills = set()
        for edge in intra_slice_edges:
            connected_skills.add(edge[0][0])  # source skill
            connected_skills.add(edge[1][0])  # target skill
        
        # Đảm bảo target_skill luôn được thêm vào (ngay cả khi isolated)
        connected_skills.add(target_skill)
        
        # Thêm inter_slice edges CHỈ cho connected skills
        inter_slice_edges = [((skill, 0), (skill, 1)) for skill in connected_skills]

        # Observation edges (chỉ nối t=0, pgmpy tự nhân bản sang t=1 khi khởi tạo lần đầu)
        observation_edges = [
            ((target_skill, 0), ('Correctness', 0)),
            ((target_skill, 0), ('ResponseTime', 0)),
        ]

        self.model.add_edges_from(intra_slice_edges + inter_slice_edges + observation_edges)

    def _add_conditional_probabilities(self, target_skill):
        self._update_initial_cpds() 

        cpds_t1 = []
        for skill in self.skills:
            parents = sorted([p for p in self.model.get_parents((skill, 1)) if p[1] == 0])
            prereq_parents_t1 = sorted([p for p in self.model.get_parents((skill, 1)) if p[1] == 1])
            evidence = [(skill, 0)] + prereq_parents_t1
            evidence_card = [2] * len(evidence)

            values = []
            for states in np.ndindex(tuple(evidence_card)):
                skill_t0_state = states[0]
                prereq_states_t1 = states[1:]
                prob_mastered_t1 = 0.0
                if skill_t0_state == 1:
                    prob_mastered_t1 = 1 - self.params['forget_rate']
                else:
                    if all(p_state == 1 for p_state in prereq_states_t1):
                        prob_mastered_t1 = self.params['learn_rate']
                    else:
                        prob_mastered_t1 = self.params['guess_rate'] * 0.1
                values.append([1 - prob_mastered_t1, prob_mastered_t1])

            values = np.array(values).T
            cpd = TabularCPD(variable=(skill, 1), variable_card=2,
                             evidence=evidence, evidence_card=evidence_card, values=values)
            cpds_t1.append(cpd)

        self._add_observation_cpds(target_skill, cpds_t1) 

    def _add_observation_cpds(self, target_skill, cpd_list=None):
        """Hàm phụ trợ để tạo và thêm CPT quan sát"""
        
        cpd_correct = TabularCPD(
            variable=('Correctness', 1), variable_card=2,
            evidence=[(target_skill, 1)], evidence_card=[2],
            values=[[1 - self.params['guess_rate'], self.params['slip_rate']],
                    [self.params['guess_rate'], 1 - self.params['slip_rate']]]
        )
        
        # GIẢM ẢNH HƯỞNG của ResponseTime - làm phân bố đều hơn
        cpd_time = TabularCPD(
            variable=('ResponseTime', 1), variable_card=3,
            evidence=[(target_skill, 1)], evidence_card=[2],
            values=[[0.5, 0.2],   # Chậm: 50% vs 20% (trước: 70% vs 10%)
                    [0.3, 0.4],   # TB: 30% vs 40% (trước: 20% vs 30%)
                    [0.2, 0.4]]   # Nhanh: 20% vs 40% (trước: 10% vs 60%)
        )
        
        cpd_correct_t0 = TabularCPD(
            variable=('Correctness', 0), variable_card=2,
            evidence=[(target_skill, 0)], evidence_card=[2],
            values=[[1 - self.params['guess_rate'], self.params['slip_rate']],
                    [self.params['guess_rate'], 1 - self.params['slip_rate']]]
        )
        cpd_time_t0 = TabularCPD(
            variable=('ResponseTime', 0), variable_card=3,
            evidence=[(target_skill, 0)], evidence_card=[2],
            values=[[0.5, 0.2], [0.3, 0.4], [0.2, 0.4]]
        )

        new_cpds = [cpd_correct, cpd_time, cpd_correct_t0, cpd_time_t0]
        
        if cpd_list is not None:
            cpd_list.extend(new_cpds)
            self.model.add_cpds(*cpd_list)
        else:
            self.model.add_cpds(*new_cpds)

    def update_belief_state(self, observation, current_target_skill=None):
        if current_target_skill is not None and current_target_skill != self.target_skill:
            print(f"Chuyển target skill từ {self.target_skill} sang {current_target_skill}")
            self._switch_target_structure(current_target_skill)

        posteriors = self.inference_engine.forward_inference(
            variables=[(skill, 1) for skill in self.skills], 
            evidence=observation
        )

        updated_beliefs = {}
        for skill in self.skills:
            posterior_cpd = posteriors[(skill, 1)]
            prob_mastered = posterior_cpd.values[1]
            updated_beliefs[skill] = prob_mastered

        self.belief_state = updated_beliefs
        self._update_initial_cpds()
        self.model.initialize_initial_state()
        self.inference_engine = DBNInference(self.model)
        
        return self.belief_state

    def _switch_target_structure(self, new_skill):
        old_skill = self.target_skill
        
        # 1. Xóa các cạnh quan sát cũ (cho cả t=0 và t=1)
        # Kiểm tra xem cạnh có tồn tại không trước khi xóa để tránh lỗi
        if self.model.has_edge((old_skill, 0), ('Correctness', 0)):
            self.model.remove_edge((old_skill, 0), ('Correctness', 0))
        if self.model.has_edge((old_skill, 0), ('ResponseTime', 0)):
            self.model.remove_edge((old_skill, 0), ('ResponseTime', 0))
            
        # Xóa cạnh ở t=1 nữa
        if self.model.has_edge((old_skill, 1), ('Correctness', 1)):
            self.model.remove_edge((old_skill, 1), ('Correctness', 1))
        if self.model.has_edge((old_skill, 1), ('ResponseTime', 1)):
            self.model.remove_edge((old_skill, 1), ('ResponseTime', 1))
        
        # Xóa CPD cũ
        cpds_to_remove = []
        for cpd in self.model.get_cpds():
            if cpd.variable[0] in ['Correctness', 'ResponseTime']:
                cpds_to_remove.append(cpd)
        for cpd in cpds_to_remove:
            self.model.remove_cpds(cpd)

        # 2. Thêm cạnh mới (cho cả t=0 và t=1)
        self.model.add_edge((new_skill, 0), ('Correctness', 0))
        self.model.add_edge((new_skill, 0), ('ResponseTime', 0))
        
        # Thêm cạnh ở t=1 để khớp với CPD
        self.model.add_edge((new_skill, 1), ('Correctness', 1))
        self.model.add_edge((new_skill, 1), ('ResponseTime', 1))
        
        # 3. Thêm CPT mới
        self._add_observation_cpds(new_skill)
        
        self.target_skill = new_skill
        
        # Bây giờ gọi hàm này sẽ an toàn
        self.model.initialize_initial_state()

    def _update_initial_cpds(self):
        cpds_to_remove = []
        for cpd in self.model.get_cpds():
            if cpd.variable[1] == 0 and cpd.variable[0] in self.skills:
                cpds_to_remove.append(cpd)
        for cpd in cpds_to_remove:
            self.model.remove_cpds(cpd)
        
        new_cpds_t0 = []
        for skill in self.skills:
            prob_mastered = self.belief_state[skill]
            prob_not_mastered = 1 - prob_mastered
            
            if not self.prerequisites.get(skill):
                cpd = TabularCPD(variable=(skill, 0), variable_card=2, 
                                values=[[prob_not_mastered], [prob_mastered]])
            else:
                prereqs = self.prerequisites[skill]
                cardinality = [2] * len(prereqs)
                n_combinations = int(np.prod(cardinality))
                prob_mastered_list = [prob_mastered] * n_combinations
                prob_not_mastered_list = [prob_not_mastered] * n_combinations
                values = np.vstack([prob_not_mastered_list, prob_mastered_list])
                
                cpd = TabularCPD(variable=(skill, 0), variable_card=2,
                                evidence=[(p, 0) for p in prereqs],
                                evidence_card=cardinality,
                                values=values)
            new_cpds_t0.append(cpd)
        
        self.model.add_cpds(*new_cpds_t0)
        self.model.check_model()

dbn_instance = None

@app.route('/', methods=['GET'])
def home():
    return jsonify({"message": "Flask API is running!", "status": "OK"}), 200

@app.route('/initialize', methods=['POST'])
def initialize():
    try:
        data = request.get_json()
        if not data:
            return jsonify({"error": "No JSON data received"}), 400

        skills = data.get('skills')
        prerequisites = data.get('prerequisites')
        params = data.get('params')
        initial_target_skill = data.get('initial_target_skill')
        belief_state = data.get('belief_state')

        if not all([skills, prerequisites is not None, params, initial_target_skill]):
            return jsonify({"error": "Missing required fields: skills, prerequisites, params, initial_target_skill"}), 400

        dbn = StudentModelDBN(skills, prerequisites, params, initial_target_skill)
        if belief_state:
            dbn._merge_external_beliefs(belief_state)
            dbn._update_initial_cpds()
            dbn.model.initialize_initial_state()
            dbn.inference_engine = DBNInference(dbn.model)

        return jsonify({
            "message": "DBN initialized successfully",
            "initial_belief_state": dbn.belief_state
        }), 200
    except Exception as e:
        return jsonify({"error": str(e)}), 500

@app.route('/update_belief_state', methods=['POST'])
def update_belief_state():
    try:
        data = request.get_json()
        print(f"[DEBUG] Received data: {data}")  # Debug log
        
        if not data:
            return jsonify({"error": "No JSON data received"}), 400

        skills = data.get('skills')
        prerequisites = data.get('prerequisites')
        params = data.get('params')
        current_target_skill = data.get('current_target_skill')
        belief_state = data.get('belief_state')
        correctness = data.get('correctness')
        response_time = data.get('response_time')
        
        print(f"[DEBUG] Skills: {skills}, Target: {current_target_skill}, Correctness: {correctness}, ResponseTime: {response_time}")

        if correctness is None or response_time is None or not all([skills, prerequisites is not None, params, current_target_skill]):
            return jsonify({"error": "Missing required fields: skills, prerequisites, params, current_target_skill, correctness, response_time"}), 400

        print(f"[DEBUG] Creating DBN model...")
        dbn = StudentModelDBN(skills, prerequisites, params, current_target_skill)
        
        if belief_state:
            print(f"[DEBUG] Merging external beliefs: {belief_state}")
            dbn._merge_external_beliefs(belief_state)
            dbn._update_initial_cpds()
            dbn.model.initialize_initial_state()
            dbn.inference_engine = DBNInference(dbn.model)

        observation = {
            ('Correctness', 1): int(correctness),
            ('ResponseTime', 1): int(response_time)
        }
        print(f"[DEBUG] Updating belief state with observation: {observation}")
        updated_belief_state = dbn.update_belief_state(observation, current_target_skill)
        
        print(f"[DEBUG] Success! Updated belief: {updated_belief_state}")
        return jsonify({
            "message": "Belief state updated successfully",
            "updated_belief_state": updated_belief_state
        }), 200
    except Exception as e:
        import traceback
        error_detail = traceback.format_exc()
        print(f"[ERROR] Exception in update_belief_state: {error_detail}")
        return jsonify({"error": str(e), "detail": error_detail}), 500

def main():
    # --- Cấu trúc kỹ năng ---
    skills_hierarchy = {
        '1': [],
        '2': [],
        '3': ['1', '2'],
    }
    ordered_skills = list(skills_hierarchy.keys())

    params_student1 = {
        'learn_rate': 0.05, 
        'forget_rate': 0.15, 
        'slip_rate': 0.25, 
        'guess_rate': 0.15
    }

    observations_sequence = [
        {'target': '3', 'obs': {('Correctness', 1): 0, ('ResponseTime', 1): 0}}, 
        {'target': '3', 'obs': {('Correctness', 1): 0, ('ResponseTime', 1): 0}}, 
        {'target': '3', 'obs': {('Correctness', 1): 0, ('ResponseTime', 1): 0}}, 
    ]

    print("\n--- Mô phỏng nhiều kỹ năng ---")
    initial_skill = observations_sequence[0]['target']
    dbn1 = StudentModelDBN(ordered_skills, skills_hierarchy, params_student1, initial_skill)
    
    print(f"\n Trạng thái niềm tin ban đầu:")
    for s, p in dbn1.belief_state.items(): print(f"  {s}: {p:.4f}")
    
    for i, item in enumerate(observations_sequence, 1):
        target = item['target']
        obs = item['obs']
        print(f"\n--- Bước {i}: Target Skill = {target} | Obs = {obs} ---")
        
        dbn1.update_belief_state(obs, current_target_skill=target)
        
        print(f"Trạng thái niềm tin sau cập nhật:")
        for s, p in dbn1.belief_state.items(): print(f"  {s}: {p:.4f}")

# Dùng cho cPanel Passenger WSGI
# Passenger sẽ tự động load biến 'application'
application = app

if __name__ == "__main__":
    # Chỉ chạy khi test local, không chạy trên production
    app.run(host="0.0.0.0", port=5000, debug=True)