# Import necessary libraries
import os
import numpy as np
import json

from typing import List, Dict, Any
from collections import defaultdict

# TODO: Load alias training files
alias_data = {}

# TODO: Load suspect data
suspect_data = {}

# TODO: Some of the activities seem related to the weapons found at the party, I wonder if we can use that to improve our decision tree.
# Note: My solution has three activities per weapon for a total of 18 pairings

activity_weapon_mapping = {
    'Ripping wippets': 'Gas', # Freebie as an example
}

# TODO: Write a function to convert the categorical data we have into numerical data for the decision tree
def extract_features(data, is_alias=True):
    rooms = ['Study', 'Dining Room', 'Kitchen', 'Pantry', 'Living Room', 'Bathroom']
    weapons = ['Poison', 'Knife', 'Gas', 'Rope', 'Bag', 'Firearm']

    features = {}

    room_time = {room: 0 for room in rooms}
    activities = []

    total_time = 0

    # TODO: Count the amount of times the suspect or alias visits each room and a list of activities they do

    # TODO: Calculate the percentage of time the suspect or alias spends in each room 

    # TODO: Create features indicating if an activity they've done suggests a weapon

    return features

import pandas as pd
from sklearn.tree import DecisionTreeClassifier
from sklearn.tree import export_text

# TODO: Prepare alias training data using the feature extraction function
alias_features_dict = {}

# TODO: Prepare suspect data using the feature extraction function
suspect_features_dict = {}

# TODO: Train Decision Tree Classifier
# I used the arguments criterion='entropy', random_state=42, max_depth=4, min_samples_split=4
# Feel free to play around with different values for the arguments

# TODO: Classify Suspects Probabilistically as each Alias

# TODO: Print your results and store them for use in the final task
decision_tree_results = {}

import numpy as np
from typing import List, Dict, Tuple
from collections import defaultdict

# TODO: Create the blueprint dictionary structure

# TODO: Initial state distribution (uniform)

# TODO: Create the emission probability matrix using objective evidence (GPS and witnesses)
def weight_room_matrix(time: str, gps_data: List[Dict], witness_statements: Dict[str, Dict]) -> np.ndarray:
    """
    Create emission probability matrix using objective evidence (GPS and witnesses).
    """

    # Start with uniform distribution
    emission = np.ones(n_states) / n_states
    
    # TODO: 1. Weight based on GPS evidence
    
    # TODO: 2. Weight based on witness evidence
    
    # TODO: Normalize Probabilities

    return emission

# TODO: Use Viterbi algorithm to find most likely sequence of true locations based on GPS and witness data
def viterbi(suspect_data: Dict, timestamps: List[str]) -> Tuple[List[str], float]:
    """
    Use Viterbi algorithm to find most likely sequence of true locations.
    Also returns the probability of the sequence.
    """
    total_timesteps = len(timestamps)
    num_rooms = len(blueprint["rooms"])
    
    # TODO: Initialize matrices
    
    
    # TODO: Create emission matrices for all timestamps
    
    # Initialize first timestep
    V[0] = np.log(initial) + np.log(emissions[0])
    
    # TODO: Finish the Viterbi forward pass
    for time_step in range(1, total_timesteps):
        for _, room_info in blueprint["rooms"].items():

            # Calculate probabilities of coming from each previous state
            probs = []
            for neighbor_room in room_info["adjacent_to"]:
                # TODO: Calculate probability of coming from this neighbor
                # TODO: Add to probabilities

            # TODO: Find most likely previous state

    # TODO: Backtrack
    
    # TODO: Reverse path and convert to room names
    
    # TODO: Calculate sequence probability
    
    return room_path, sequence_prob
    

# TODO: Analyze each suspect's path to determine how well it matches the most likely path from the HMM

movement_results = {}

print("\nHMM Path Analysis:")
print("=" * 50)

for suspect_name, suspect in suspect_data.items():
    # Get all timestamps in chronological order
    timestamps = []
    for entry in suspect['suspect_statement']:
        timestamps.append(entry['time'])
    timestamps.sort()
    
    # Run Viterbi to get most likely true path using GPS data and witness statements
    true_path, path_prob = viterbi(
        {
            'gps_data': suspect['gps_data'],
            'witness_statements': suspect['witness_statements']
        }, 
        timestamps
    )
    
    # TODO: Get suspect's stated path
    
    # TODO: Calculate matches between true and stated paths
    
    print(f"\n{suspect_name}:")
    print(f'\tAgreement of Statement Path with HMM Path: {agreement}, Number of Contradictions: {len(true_path) - matches}')

    path_agreement = {
        'path_probability': path_prob,
        'statement_agreement': agreement,
        'num_contradictions': len(true_path) - matches
    }
    
    movement_results[suspect_name] = path_agreement

def combine_evidence_bayesian(aliases, decision_tree_results, hmm_results) -> Dict[str, Dict]:
    # TODO: Combine all evidence, assuming that whomever is Mr. Green is lying!
    # TODO: P(Alias|Evidence) ∝ P(Alias|Activities) * P(Alias|Truthfulness)

    final_results = {}
    
    # We can assume that whomever is Mr. Green is lying
    # So low statement match with highest likelihood path → higher probability of being Mr. Green

    for suspect_name in decision_tree_results.keys():
        # TODO: Get decision tree probabilities
        dt_probs = decision_tree_results[suspect_name]
        
        # TODO: Get path agreement from the HMM results 
        path_agreement = hmm_results[suspect_name]['statement_agreement']
        
        # TODO: Calculate posterior probabilities
        posterior = {}
        
        # TODO: Normalize
        
        final_results[suspect_name] = { alias: prob/total for alias, prob in posterior.items() }
        
    return final_results

# Combine evidence using Bayesian reasoning
final_results = combine_evidence_bayesian(aliases, decision_tree_results, movement_results)

for suspect, results in final_results.items():
    print(f"{suspect}:")
    for alias, prob in results.items():
        print(f"  {alias}: {prob:.3f}")