import gradio as gr
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation, PillowWriter
import scipy.io.wavfile as wavfile

# Constants for sound generation
SAMPLE_RATE = 48000
COLUMN_DURATION = 1  # Duration of each column in seconds

# Mapping of matrix numbers to musical notes
notes = ["A", "A#", "B", "C", "C#", "D", "D#", "E", "F", "F#", "G", "G#"]
note_map = {
    2: "A", 3: "A#", 4: "B", 5: "C", 6: "C#", 7: "D", 8: "D#", 9: "E", 
    10: "F", 11: "F#", 12: "G", 13: "G#"
}

def generate_tone(note, duration):
    """Generates a tone for a specified note and duration."""
    frequency = 440 * 2 ** ((notes.index(note) - 9) / 12)
    t = np.linspace(0, duration, int(SAMPLE_RATE * duration), endpoint=False)
    return 0.5 * np.sin(2 * np.pi * frequency * t)

def get_matrix_audio(matrix):
    """Generate audio sequence for one matrix state."""
    audio_sequence = []
    
    for col in range(matrix.shape[1]):
        col_audio = np.zeros(0)
        for row in range(matrix.shape[0]):
            num = matrix[row, col]
            if num > 1:
                note = note_map.get(num, "A")
                tone = generate_tone(note, duration=COLUMN_DURATION / matrix.shape[0])
                col_audio = np.concatenate((col_audio, tone))
        
        audio_sequence.append(col_audio)
    
    return np.concatenate(audio_sequence)

class NumberSpreadSimulator:
    def __init__(self, initial_matrix):
        self.grid = np.array(initial_matrix, dtype=int)
        self.audio_frames = []
        self.original_numbers = self.find_original_numbers()
        self.initialize_audio()
        
    def find_original_numbers(self):
        """Find all original numbers in the initial matrix."""
        numbers = set()
        for row in range(self.grid.shape[0]):
            for col in range(self.grid.shape[1]):
                if self.grid[row, col] > 1:
                    numbers.add(self.grid[row, col])
        return list(numbers)
    
    def initialize_audio(self):
        """Generate audio for initial state and store it."""
        self.audio_frames.append(get_matrix_audio(self.grid))

    def find_original_number(self, current):
        """Find the original source number based on the current number."""
        closest = min(self.original_numbers, key=lambda x: abs(x - current))
        return closest
    
    def step(self):
        """Simulates a step in the matrix spread."""
        new_grid = np.zeros_like(self.grid)
        has_changes = False
        
        # Process each position in the grid
        for row in range(self.grid.shape[0]):
            for col in range(self.grid.shape[1]):
                current = self.grid[row, col]
                
                if current > 1:
                    has_changes = True
                    half = current // 2
                    
                    # Define potential target positions
                    targets = []
                    if current % 2 == 0:  # Even
                        if row > 0: targets.append((row - 1, col, half))
                        if col > 0: targets.append((row, col - 1, half))
                    else:  # Odd
                        if col > 0: targets.append((row, col - 1, half))
                        if row > 0: targets.append((row - 1, col, half))
                        if row > 0 and col > 0: targets.append((row - 1, col - 1, 1))
                    
                    # Process each target position
                    for target_row, target_col, value in targets:
                        if self.grid[target_row, target_col] == -1:
                            # If target is -1, replace with original number
                            new_grid[target_row, target_col] = self.find_original_number(current)
                        elif self.grid[target_row, target_col] == -2:
                            # If target is -2, double the incoming value
                            new_grid[target_row, target_col] = value * 2
                        elif self.grid[target_row, target_col] == -3:
                            # If target is -3, set to 0
                            new_grid[target_row, target_col] = 0
                        else:
                            new_grid[target_row, target_col] += value
                
                # Copy over any remaining special values
                if (self.grid[row, col] in [-1, -2, -3]) and new_grid[row, col] == 0:
                    new_grid[row, col] = self.grid[row, col]
        
        self.grid = new_grid
        if has_changes:
            self.audio_frames.append(get_matrix_audio(self.grid))
        
        return has_changes, self.grid

def create_animation(matrix):
    # Initialize simulator
    sim = NumberSpreadSimulator(matrix)
    
    # Determine matrix dimensions
    rows, cols = sim.grid.shape
    
    # Create figure and axis with dynamic limits
    fig, ax = plt.subplots(figsize=(10, 10))
    ax.set_xlim(-0.5, cols - 0.5)
    ax.set_ylim(-0.5, rows - 0.5)
    ax.set_title("Matrix Spread Visualization")
    
    # Initialize plot elements
    circles = []
    labels = []
    
    for i in range(rows):
        for j in range(cols):
            circle = plt.Circle((j, rows - 1 - i), 0.2, 
                              color='blue', 
                              fill=False)
            label = ax.text(j, rows - 1 - i, '', ha='center', va='center', fontsize=12)
            circles.append(circle)
            labels.append(label)
            ax.add_patch(circle)
    
    current_frame = [0]
    
    def update(frame):
        if current_frame[0] == 0:
            matrix = sim.grid
        else:
            has_changes, matrix = sim.step()
            if not has_changes:
                ani.event_source.stop()
                return circles + labels
        
        for i in range(rows):
            for j in range(cols):
                value = matrix[i, j]
                index = i * cols + j
                
                if value != 0:
                    circles[index].set_radius(0.1 + 0.1 * (abs(value) / 10))
                    if value == -1:
                        circles[index].set_facecolor('green')
                    elif value == -2:
                        circles[index].set_facecolor('purple')
                    elif value == -3:
                        circles[index].set_facecolor('red')
                    else:
                        circles[index].set_facecolor('orange')
                else:
                    circles[index].set_radius(0.1)
                    circles[index].set_facecolor('blue')
                labels[index].set_text(str(value) if value != 0 else '')
        
        current_frame[0] += 1
        return circles + labels
    
    # Estimate the number of frames needed
    max_steps = 100  # Set a reasonable limit for the number of steps
    ani = FuncAnimation(fig, update, frames=max_steps, interval=1000, blit=True, cache_frame_data=False)
    
    # Save the animation to a GIF file
    gif_path = "matrix_animation.gif"
    ani.save(gif_path, writer=PillowWriter(fps=1))
    plt.close(fig)
    
    return gif_path

def run_simulation(matrix_input):
    """
    Run the full simulation based on user-input matrix
    
    :param matrix_input: 2D list of integers representing the matrix
    :return: tuple of (audio_path, gif_path)
    """
    # Convert input to numpy array
    matrix = np.array(matrix_input, dtype=int)
    
    # Initialize simulator
    sim = NumberSpreadSimulator(matrix)
    
    # Run simulation until no more changes
    while True:
        has_changes, _ = sim.step()
        if not has_changes:
            break
    
    # Generate audio
    final_audio = np.concatenate(sim.audio_frames)
    final_audio = np.int16(final_audio * 32767)
    
    # Save audio
    audio_path = "matrix_sound.wav"
    wavfile.write(audio_path, SAMPLE_RATE, final_audio)
    
    # Create animation
    gif_path = create_animation(matrix)
    
    return audio_path, gif_path

# Gradio Interface
def create_gradio_interface():
    # Default initial matrix
    default_matrix = [
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 47]
    ]

    # Create Gradio interface
    iface = gr.Interface(
        fn=run_simulation,
        inputs=[
            gr.Dataframe(
                headers=[str(i) for i in range(10)],
                datatype="number",
                value=default_matrix,
                type="numpy",
                label="Edit Matrix Values"
            )
        ],
        outputs=[
            gr.Audio(type="filepath", label="Generated Sound"),
            gr.Image(type="filepath", label="Matrix Animation")
        ],
        title="Number Spread Simulator",
        description="Edit the matrix and see how numbers spread, generating a unique sound and animation!"
    )
    
    return iface

# Launch the interface
iface = create_gradio_interface()

if __name__ == "__main__":
    iface.launch()