import numpy as np
import matplotlib.pyplot as plt

class DiscreteBayesFilter1D:
    def __init__(self, size, door_positions):
        """
        TASK 1A: Initializing the Environment and Beliefs
        size: int - number of cells in the corridor
        door_positions: list - indexes of cells where there are doors
        """
        self.size = size
        
        # Creating a map (1 where the door is, 0 everywhere else)
        self.map = np.zeros(size)
        self.map[door_positions] = 1
        
        # TODO: Initialize 'self.belief' as a uniform distribution.
        # Remember: the sum of the probabilities in the array must be exactly 1.0!
        # Hint: use np.ones() or np.full() and divide by size.
        self.belief = np.zeros(size) # Change it!


    def plot_belief(self, title="Robot Belief"):
        """
        TASK 1B: Visualization
        Create a bar chart showing the current probability distribution.
        """
        plt.figure(figsize=(12, 4))
        
        # TODO: Use plt.bar() to plot probability bars for each cell.
        # x = range from 0 to self.size, y = values ​​from self.belief
        
        # TODO: Add a guide line (e.g. plt.step) that shows the location of the door from the map.
        # This will make it easier to interpret whether the robot "knows" it is at the door.

        plt.title(title)
        plt.xlabel("Position (cell index)")
        plt.ylabel("Probability P(x)")
        plt.ylim(0, 1.1) # Fixed scale makes it easier to compare graphs
        plt.grid(axis='y', alpha=0.3)
        plt.show()

    def predict(self, move, kernel):
            """
            TASK 2: Implementing the traffic model
            move: int - how many squares the robot intends to move
            kernel: np.array - error kernel (e.g. [0.1, 0.8, 0.1])
            """
            # TODO: 1. Deterministic Shift
            # Move values ​​in self.belief by 'move' position.
            # Hint: np.roll(array, shift) works great for cyclic maps.
            prior_belief = np.zeros(self.size) 
            
            # TODO: 2. Convolution
            # Apply kernel on the shifted belief to add process noise.
            # Use np.convolve(array, kernel, mode='same').
            # Write the result back to self.belief.
            
            # self.belief = ...
            pass

if __name__ == "__main__":
    # Test parameters
    N = 20
    doors = [3, 4, 10, 15]
    
    # Creating a filter instance
    robot_filter = DiscreteBayesFilter1D(size=N, door_positions=doors)
    
    # Displaying the initial state
    print(f"Sum of probabilities: {np.sum(robot_filter.belief)}")
    robot_filter.plot_belief("Initial state - the robot does not know where it is")