"""
logos/baker.py - The SPCW Encoder (Fractal Compression)
"""

import numpy as np
from .fractal_engine import ContextAction

class BreadBaker:
    """
    Encodes image tiles into Prime Atoms using Quadtree Decomposition.
    """
    
    def __init__(self, variance_threshold=30.0, min_block_size=8):
        self.variance_threshold = variance_threshold
        self.min_block_size = min_block_size
    
    def bake(self, tile: np.ndarray, prefix_path: list = None) -> list:
        """
        Compress a tile into atom definitions.
        Returns list of dicts: {'path_bits': [], 'payload': bytes}
        """
        self.atoms = []
        if prefix_path is None:
            prefix_path = []
            
        self._recursive_decompose(tile, level=0, path_bits=prefix_path)
        return self.atoms

    def _recursive_decompose(self, block: np.ndarray, level: int, path_bits: list):
        h, w = block.shape[:2]
        if h == 0 or w == 0: return

        mean_val = np.mean(block, axis=(0, 1))
        variance = np.var(block)
        
        # Stop condition: Low variance OR too small OR max depth
        # Max local depth 5 (512 -> 16px blocks)
        is_leaf = (variance < self.variance_threshold) or (w <= self.min_block_size) or (level >= 5)
        
        if is_leaf:
            self._emit_atom(mean_val, path_bits)
        else:
            # Subdivide
            half_h, half_w = h // 2, w // 2
            # TL, TR, BL, BR
            self._recursive_decompose(block[:half_h, :half_w], level + 1, path_bits + [0b00])
            self._recursive_decompose(block[:half_h, half_w:], level + 1, path_bits + [0b01])
            self._recursive_decompose(block[half_h:, :half_w], level + 1, path_bits + [0b10])
            self._recursive_decompose(block[half_h:, half_w:], level + 1, path_bits + [0b11])

    def _emit_atom(self, rgb_mean, path_bits):
        r, g, b = int(rgb_mean[0]), int(rgb_mean[1]), int(rgb_mean[2])
        matrix_key = (b << 16) | (g << 8) | r 
        
        # --- ATOM 1: WARM UP (Harmonic Resonance) ---
        # We must raise heat_state from 0 to 1 so that GF4 multiplication (Dissolution)
        # has a non-zero scalar. 0 * Input = 0 (No Change).
        # ContextAction.PERSIST (00) increments heat.
        
        warm_action = ContextAction.PERSIST
        warm_ctrl = (warm_action.value << 6) & 0xC0
        warm_payload = bytes([warm_ctrl])
        
        self.atoms.append({
            "path_bits": path_bits,
            "payload": warm_payload
        })
        
        # --- ATOM 2: STRUCTURE (Matrix Key) ---
        # Pack payload compatible with LogoFractalEngine vector decoder
        # We pack matrix_key (RGB) into the vector slots
        action = ContextAction.CHANGE_1
        control_byte = (action.value << 6) & 0xC0
        
        payload_bytes = bytearray()
        payload_bytes.append(control_byte)
        
        # Convert RGB integer to 2-bit vector sequence
        vector = [(matrix_key >> (i*2)) & 0b11 for i in range(16)]
        
        # Pack vector into 4 bytes
        # Note: reconstruction expects this exact packing to retrieve matrix_key
        # Byte order matters. We pack vector[0..3] into first byte, etc.
        # This aligns with `_payload_to_vector` in fractal_engine.
        
        # Re-verify _payload_to_vector in fractal_engine:
        # for byte in payload[:8]:
        #    append((byte >> 6) & 0b11) ...
        
        # So we pack vector[0] into MSB of byte 0?
        # No, `_payload_to_vector` extracts MSB first: `(byte >> 6) & 0b11`.
        # So we should pack vector[0] at MSB.
        
        current_vector_idx = 0
        for _ in range(4): # 4 bytes to hold 16 pairs
            b_val = 0
            if current_vector_idx < 16: b_val |= (vector[current_vector_idx] & 0b11) << 6
            if current_vector_idx+1 < 16: b_val |= (vector[current_vector_idx+1] & 0b11) << 4
            if current_vector_idx+2 < 16: b_val |= (vector[current_vector_idx+2] & 0b11) << 2
            if current_vector_idx+3 < 16: b_val |= (vector[current_vector_idx+3] & 0b11)
            
            payload_bytes.append(b_val)
            current_vector_idx += 4
            
        self.atoms.append({
            "path_bits": path_bits,
            "payload": bytes(payload_bytes)
        })