import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer
from dataclasses import dataclass
import copy
import asyncio
import uuid
import time

# Helper to allow dot-notation access (chunk.choices[0].delta.content)
class OpenAIObject(dict):
    def __getattr__(self, name):
        if name in self:
            value = self[name]
            if isinstance(value, dict):
                return OpenAIObject(value)
            if isinstance(value, list):
                return [OpenAIObject(v) if isinstance(v, dict) else v for v in value]
            return value
        raise AttributeError(f"'OpenAIObject' object has no attribute '{name}'")
    
    def __setattr__(self, name, value):
        self[name] = value

@dataclass
class ZIPRCConfig:
    model_name: str = "dataopsnick/Qwen3-4B-Instruct-2507-zip-rc"
    reward_bins: int = 8
    length_bins: int = 7
    total_zip_tokens: int = 56
    zip_start_offset: int = 56
    alpha: float = 0.1      
    beta: float = 0.05      
    smoothing_window: int = 3
    r_boundaries = torch.linspace(0, 1, 9)
    l_boundaries = torch.tensor([0, 16, 32, 64, 128, 256, 512, 1024], dtype=torch.float32)

class ZIPRCMath:
    @staticmethod
    def get_bin_idx(val, boundaries):
        for i in range(len(boundaries) - 1):
            if boundaries[i] <= val < boundaries[i+1]:
                return i
        return len(boundaries) - 2

    @staticmethod
    def apply_horizon_capping(joint_probs, current_len, horizon, config):
        """Eq 25: Collapses mass where length > horizon into a failure state."""
        B, R_bins, L_bins = joint_probs.shape
        device = joint_probs.device
        cutoff_l_idx = L_bins - 1
        for i, bound in enumerate(config.l_boundaries):
            if bound > horizon:
                cutoff_l_idx = max(0, i - 1)
                break

        capped_probs = joint_probs.clone()
        valid_mask = torch.zeros((L_bins), dtype=torch.bool, device=device)
        valid_mask[:cutoff_l_idx+1] = True
        kept_mass = capped_probs[:, :, valid_mask].sum(dim=(1, 2)) 
        pruned_mass = 1.0 - kept_mass
        capped_probs[:, :, ~valid_mask] = 0.0
        capped_probs[:, 0, cutoff_l_idx] += pruned_mass
        return capped_probs

    @staticmethod
    def get_marginals(joint_probs):
        q_v = joint_probs.sum(dim=2) 
        q_l = joint_probs.sum(dim=1) 
        return q_v, q_l

    @staticmethod
    def compute_expected_max_value(marginals_list, values_per_bin):
        if not marginals_list: return 0.0
        stacked_marginals = torch.cat(marginals_list, dim=0) 
        cdfs = torch.cumsum(stacked_marginals, dim=1)
        f_max = torch.prod(cdfs, dim=0) 
        f_max_shifted = torch.roll(f_max, 1)
        f_max_shifted[0] = 0.0
        p_max = f_max - f_max_shifted
        expected_max = torch.sum(p_max * values_per_bin).item()
        return expected_max

    @staticmethod
    def compute_sampling_utility(candidates, config):
        """Eq 19: Utility optimization for shared horizon."""
        if not candidates: return -1e9
        device = candidates[0]['joint_probs'].device
        r_vals = (config.r_boundaries[:-1] + config.r_boundaries[1:]).to(device) / 2
        l_vals = (config.l_boundaries[:-1] + config.l_boundaries[1:]).to(device) / 2

        sum_est_total_len = 0.0
        for cand in candidates:
            qv, ql = ZIPRCMath.get_marginals(cand['joint_probs'].unsqueeze(0))
            e_rem_len = torch.sum(ql * l_vals).item()
            curr_prefix_len = cand['ids'].shape[1]
            sum_est_total_len += (curr_prefix_len + e_rem_len)

        b_bar = max(sum_est_total_len / len(candidates) if len(candidates) > 0 else 1.0, 1.0)
        beta_tilde = config.beta / b_bar

        best_util = -float('inf')
        search_space = config.l_boundaries.tolist() + [2048]
        for h in search_space:
            h = int(h)
            q_v_list, q_l_list = [], []
            for cand in candidates:
                capped_joint = ZIPRCMath.apply_horizon_capping(
                    cand['joint_probs'].unsqueeze(0), cand['current_len'], h, config
                )
                qv, ql = ZIPRCMath.get_marginals(capped_joint)
                q_v_list.append(qv)
                q_l_list.append(ql)
            e_max_reward = ZIPRCMath.compute_expected_max_value(q_v_list, r_vals)
            e_latency = ZIPRCMath.compute_expected_max_value(q_l_list, l_vals)
            total_compute = sum(torch.sum(ql * l_vals).item() for ql in q_l_list)
            cost = beta_tilde * (config.alpha * total_compute + (1 - config.alpha) * e_latency)
            util = e_max_reward - cost
            if util > best_util: best_util = util
        return best_util

class PredictionBuffer:
    def __init__(self, window_size):
        self.window = window_size
        self.history = []
    def add(self, prob_tensor):
        self.history.append(prob_tensor)
        if len(self.history) > self.window: self.history.pop(0)
    def get_smoothed(self):
        stack = torch.stack(self.history)
        return torch.mean(stack, dim=0)

class ZIPRCModel(torch.nn.Module):
    def __init__(self, config: ZIPRCConfig):
        super().__init__()
        self.config = config
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.base_model = AutoModelForCausalLM.from_pretrained(
            config.model_name, torch_dtype=torch.bfloat16, device_map=self.device
        )
        self.tokenizer = AutoTokenizer.from_pretrained(config.model_name)
        self.zip_start_id = self.base_model.config.vocab_size - config.zip_start_offset
        self.base_model.eval()

    def get_joint_distribution(self, logits):
        zip_logits = logits[:, self.zip_start_id : self.zip_start_id + self.config.total_zip_tokens]
        probs = F.softmax(zip_logits, dim=-1)
        return probs.view(-1, self.config.reward_bins, self.config.length_bins)

class ZIPRCSampler:
    def __init__(self, model):
        self.model = model
        self.config = model.config

    def select_best_trajectory(self, trajectories):
        if not trajectories: return None
        best_traj = None
        best_score = -float('inf')
        device = trajectories[0]['joint_probs'].device
        r_vals = (self.config.r_boundaries[:-1] + self.config.r_boundaries[1:]).to(device) / 2
        for traj in trajectories:
            qv, _ = ZIPRCMath.get_marginals(traj['joint_probs'].unsqueeze(0))
            score = torch.sum(qv * r_vals).item()
            traj['final_score'] = score
            if score > best_score:
                best_score = score
                best_traj = traj
        return best_traj

    async def openai(self, messages, max_tokens=512, initial_samples=2):
        """
        Async generator that yields OpenAI-compatible chunks with added ZIP-RC introspection data.
        """
        # 1. Handle Input (String or Messages List)
        if isinstance(messages, str):
            prompt = messages
        else:
            # Assumes generic chat template
            prompt = self.model.tokenizer.apply_chat_template(
                messages, tokenize=False, add_generation_prompt=True
            )

        # 2. Setup Candidates
        input_ids = self.model.tokenizer(prompt, return_tensors="pt").input_ids.to(self.model.device)
        
        candidates = []
        for i in range(initial_samples):
            candidates.append({
                'id': i, 'ids': input_ids.clone(), 'finished': False,
                'buffer': PredictionBuffer(self.config.smoothing_window),
                'joint_probs': None, 'current_len': 0
            })
        
        finished_trajectories = []
        
        chat_id = f"chatcmpl-{uuid.uuid4()}"
        created_ts = int(time.time())
        model_name = self.config.model_name

        # State for delta streaming
        last_top1_id = -1
        last_top1_len = input_ids.shape[1]

        for step in range(max_tokens):
            if not candidates: break

            # --- [A] MODEL FORWARD (Async Wrapper) ---
            active_ids = torch.cat([c['ids'] for c in candidates], dim=0)
            
            # Simple wrapper to allow event loop to breathe
            await asyncio.sleep(0) 
            
            with torch.no_grad():
                outputs = self.model.base_model(active_ids)
                next_token_logits = outputs.logits[:, -1, :]
                raw_joint = self.model.get_joint_distribution(next_token_logits)

            # --- Update Candidates ---
            for i, c in enumerate(candidates):
                c['buffer'].add(raw_joint[i])
                c['joint_probs'] = c['buffer'].get_smoothed()
                c['current_len'] = step
                
                valid_logits = next_token_logits[i].clone()
                valid_logits[self.model.zip_start_id : self.model.zip_start_id + self.config.total_zip_tokens] = -float('inf')
                probs = F.softmax(valid_logits, dim=-1)
                next_token = torch.multinomial(probs, 1).unsqueeze(0)
                
                c['ids'] = torch.cat([c['ids'], next_token], dim=1)
                if next_token.item() == self.model.tokenizer.eos_token_id:
                    c['finished'] = True
                    finished_trajectories.append(c)

            candidates = [c for c in candidates if not c['finished']]
            if not candidates: break

            # --- [B] META-ACTIONS ---
            cand_metrics = []
            r_vals = (self.config.r_boundaries[:-1] + self.config.r_boundaries[1:]).to(self.model.device)/2
            for i, c in enumerate(candidates):
                qv, _ = ZIPRCMath.get_marginals(c['joint_probs'].unsqueeze(0))
                e_r = torch.sum(qv * r_vals).item()
                cand_metrics.append((i, e_r))
            
            sorted_by_reward = sorted(cand_metrics, key=lambda x: x[1], reverse=True)
            top_indices = [x[0] for x in sorted_by_reward]
            
            possible_actions = [('keep', candidates)]
            MAX_SAMPLES = 8

            if len(candidates) < MAX_SAMPLES:
                top_idx = top_indices[0]
                new_set = copy.deepcopy(candidates)
                clone = copy.deepcopy(new_set[top_idx])
                clone['id'] = max([c['id'] for c in new_set], default=0) + 1
                new_set.append(clone)
                possible_actions.append(('branch_top1', new_set))
                
                if len(candidates) >= 2 and len(candidates) + 1 < MAX_SAMPLES:
                    new_set_b2 = copy.deepcopy(candidates)
                    clone2 = copy.deepcopy(new_set_b2[top_indices[1]])
                    clone2['id'] = max([c['id'] for c in new_set_b2], default=0) + 1
                    new_set_b2.append(clone2)
                    possible_actions.append(('branch_top2', new_set_b2))

            if len(candidates) > 1:
                worst_idx = top_indices[-1]
                new_set = [c for i, c in enumerate(candidates) if i != worst_idx]
                possible_actions.append(('prune_bot1', new_set))

            if len(candidates) > 1 and top_indices[0] != top_indices[-1]:
                top_id = candidates[top_indices[0]]['id']
                worst_idx = top_indices[-1]
                new_set = copy.deepcopy(candidates)
                new_set = [c for i, c in enumerate(new_set) if i != worst_idx]
                source = next(c for c in new_set if c['id'] == top_id)
                clone = copy.deepcopy(source)
                clone['id'] = max([c['id'] for c in new_set], default=0) + 1
                new_set.append(clone)
                possible_actions.append(('swap', new_set))

            best_action_name, best_util, best_next_candidates = 'keep', -float('inf'), candidates
            for name, cand_set in possible_actions:
                if not cand_set: continue
                penalty = 0.0 if name == 'keep' else 0.01
                util = ZIPRCMath.compute_sampling_utility(cand_set, self.config) - penalty
                if util > best_util:
                    best_util, best_action_name, best_next_candidates = util, name, cand_set
            
            candidates = best_next_candidates

            # --- [C] PREPARE INTROSPECTION PAYLOAD ---
            vis_metrics = []
            for c in candidates:
                qv, _ = ZIPRCMath.get_marginals(c['joint_probs'].unsqueeze(0))
                e_r = torch.sum(qv * r_vals).item()
                vis_metrics.append((c, e_r))
            vis_sorted = sorted(vis_metrics, key=lambda x: x[1], reverse=True)
            
            lhs_c = vis_sorted[0][0] if len(vis_sorted) > 0 else None
            rhs_c = vis_sorted[1][0] if len(vis_sorted) > 1 else None
            lhs_score = vis_sorted[0][1] if len(vis_sorted) > 0 else 0.0
            rhs_score = vis_sorted[1][1] if len(vis_sorted) > 1 else 0.0

            def get_text(c_obj):
                if not c_obj: return ""
                curr_ids = c_obj['ids'][0]
                full_text = self.model.tokenizer.decode(curr_ids, skip_special_tokens=True)
                return full_text

            lhs_text = get_text(lhs_c)
            rhs_text = get_text(rhs_c)

            delta_content = ""
            if lhs_c and lhs_c['id'] == last_top1_id:
                 new_len = len(lhs_text)
                 if new_len > last_top1_len:
                     delta_content = lhs_text[last_top1_len:]
                     last_top1_len = new_len
            elif lhs_c:
                last_top1_id = lhs_c['id']
                last_top1_len = len(lhs_text)
                delta_content = "" 

            chunk_dict = {
                "id": chat_id,
                "object": "chat.completion.chunk",
                "created": created_ts,
                "model": model_name,
                "choices": [{"index": 0, "delta": {"content": delta_content}, "finish_reason": None}],
                "zip_rc": {
                    "step": step,
                    "action": best_action_name,
                    "utility": best_util,
                    "lhs_text": lhs_text,
                    "rhs_text": rhs_text,
                    "lhs_score": lhs_score,
                    "rhs_score": rhs_score,
                    "lhs_id": lhs_c['id'] if lhs_c else -1,
                    "rhs_id": rhs_c['id'] if rhs_c else -1
                }
            }
            yield OpenAIObject(chunk_dict)

        # Calculate Final Best Answer (clean from swaps/backtracks)
        # Include running candidates in case max_tokens was hit before EOS
        all_trajs = finished_trajectories + candidates
        best_traj = self.select_best_trajectory(all_trajs)
        final_answer = ""
        if best_traj:
             # Decode only the generated response (exclude prompt)
             prompt_len = input_ids.shape[1]
             final_ids = best_traj['ids'][0][prompt_len:]
             final_answer = self.model.tokenizer.decode(final_ids, skip_special_tokens=True)

        yield OpenAIObject({
            "id": chat_id,
            "object": "chat.completion.chunk",
            "created": created_ts,
            "model": model_name,
            "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
            "zip_rc": {
                "action": "finished",
                "final_text": final_answer
            }
        })

    def generate_stream(self, prompt, max_new_tokens=512, initial_samples=2):
        # Setup
        input_ids = self.model.tokenizer.apply_chat_template(
            [{"role": "user", "content": prompt}], tokenize=True, return_tensors="pt", add_generation_prompt=True
        ).to(self.model.device)

        candidates = []
        for i in range(initial_samples):
            candidates.append({
                'id': i, 'ids': input_ids.clone(), 'finished': False,
                'buffer': PredictionBuffer(self.config.smoothing_window),
                'joint_probs': None, 'current_len': 0
            })
        
        finished_trajectories = []
        text_cache = {} 
        
        # UI State
        dashboard_widget = None
        last_cli_height = 0
        
        # Check environment for widget vs CLI
        try:
            import ipywidgets as widgets
            from IPython.display import display
            ENV_MODE = 'notebook'
        except ImportError:
            ENV_MODE = 'cli'

        if ENV_MODE == 'notebook':
            import html
            from tqdm.notebook import tqdm
            dashboard_widget = widgets.HTML(value="Initialization...")
            display(dashboard_widget)
            pbar = tqdm(total=max_new_tokens, display=False) 
        else:
            import sys, shutil, textwrap
            from tqdm import tqdm
            pbar = tqdm(total=max_new_tokens, dynamic_ncols=True, bar_format='{bar}| {n_fmt}/{total_fmt}')

        try:
            for step in range(max_new_tokens):
                if not candidates: break

                # --- [A] MODEL FORWARD ---
                active_ids = torch.cat([c['ids'] for c in candidates], dim=0)
                with torch.no_grad():
                    outputs = self.model.base_model(active_ids)
                    next_token_logits = outputs.logits[:, -1, :]
                    raw_joint = self.model.get_joint_distribution(next_token_logits)

                for i, c in enumerate(candidates):
                    c['buffer'].add(raw_joint[i])
                    c['joint_probs'] = c['buffer'].get_smoothed()
                    c['current_len'] = step
                    
                    valid_logits = next_token_logits[i].clone()
                    valid_logits[self.model.zip_start_id : self.model.zip_start_id + self.config.total_zip_tokens] = -float('inf')
                    probs = F.softmax(valid_logits, dim=-1)
                    next_token = torch.multinomial(probs, 1).unsqueeze(0)
                    
                    c['ids'] = torch.cat([c['ids'], next_token], dim=1)
                    if next_token.item() == self.model.tokenizer.eos_token_id:
                        c['finished'] = True
                        finished_trajectories.append(c)

                candidates = [c for c in candidates if not c['finished']]
                if not candidates: break

                # --- [B] META-ACTIONS (Restored Full Logic) ---
                
                # 1. Metric Calculation & Sorting
                cand_metrics = []
                r_vals = (self.config.r_boundaries[:-1] + self.config.r_boundaries[1:]).to(self.model.device)/2
                for i, c in enumerate(candidates):
                    qv, _ = ZIPRCMath.get_marginals(c['joint_probs'].unsqueeze(0))
                    e_r = torch.sum(qv * r_vals).item()
                    cand_metrics.append((i, e_r))
                
                # Sort to identify Top-1, Top-2, and Worst
                sorted_by_reward = sorted(cand_metrics, key=lambda x: x[1], reverse=True)
                top_indices = [x[0] for x in sorted_by_reward]
                
                # 2. Define Possible Actions
                possible_actions = [('keep', candidates)]
                MAX_SAMPLES = 8

                # Action: Branch Top-1
                if len(candidates) < MAX_SAMPLES:
                    top_idx = top_indices[0]
                    new_set = copy.deepcopy(candidates)
                    clone = copy.deepcopy(new_set[top_idx])
                    clone['id'] = max([c['id'] for c in new_set], default=0) + 1
                    new_set.append(clone)
                    possible_actions.append(('branch_top1', new_set))
                    
                    # Action: Branch Top-2
                    if len(candidates) >= 2 and len(candidates) + 1 < MAX_SAMPLES:
                        new_set2 = copy.deepcopy(new_set) # Base off the set that already branched top1? No, independent action in original code.
                        # Original code treats them as distinct alternative meta-actions for the step.
                        # Re-building from clean candidates for branch_top2:
                        new_set_b2 = copy.deepcopy(candidates)
                        clone2 = copy.deepcopy(new_set_b2[top_indices[1]])
                        clone2['id'] = max([c['id'] for c in new_set_b2], default=0) + 1
                        new_set_b2.append(clone2)
                        possible_actions.append(('branch_top2', new_set_b2))

                # Action: Prune Worst 1
                if len(candidates) > 1:
                    worst_idx = top_indices[-1]
                    new_set = [c for i, c in enumerate(candidates) if i != worst_idx]
                    possible_actions.append(('prune_bot1', new_set))

                # Action: Prune Worst 2
                if len(candidates) > 2:
                    worst_indices = set(top_indices[-2:])
                    new_set = [c for i, c in enumerate(candidates) if i not in worst_indices]
                    possible_actions.append(('prune_bot2', new_set))
                
                # Action: Swap (Prune Worst, Branch Best)
                if len(candidates) > 1 and top_indices[0] != top_indices[-1]:
                    top_id = candidates[top_indices[0]]['id']
                    worst_idx = top_indices[-1]
                    new_set = copy.deepcopy(candidates)
                    new_set = [c for i, c in enumerate(new_set) if i != worst_idx]
                    source = next(c for c in new_set if c['id'] == top_id)
                    clone = copy.deepcopy(source)
                    clone['id'] = max([c['id'] for c in new_set], default=0) + 1
                    new_set.append(clone)
                    possible_actions.append(('swap', new_set))

                # 3. Select Best Action via Utility
                best_action_name, best_util, best_next_candidates = 'keep', -float('inf'), candidates
                for name, cand_set in possible_actions:
                    if not cand_set: continue
                    penalty = 0.0 if name == 'keep' else 0.01
                    util = ZIPRCMath.compute_sampling_utility(cand_set, self.config) - penalty
                    if util > best_util:
                        best_util, best_action_name, best_next_candidates = util, name, cand_set
                
                # Apply selection
                candidates = best_next_candidates

                # Re-evaluate metrics for visualization (indices might have shifted or sizes changed)
                # We need to find the new Top-1 and Top-2 to display in the UI.
                vis_metrics = []
                for i, c in enumerate(candidates):
                    qv, _ = ZIPRCMath.get_marginals(c['joint_probs'].unsqueeze(0))
                    e_r = torch.sum(qv * r_vals).item()
                    vis_metrics.append((c, e_r))
                vis_sorted = sorted(vis_metrics, key=lambda x: x[1], reverse=True)

                # --- [C] ADVANCED DASHBOARD RENDERING ---
                pbar.update(1)
                
                # 1. Prepare Text (Always show current best 2)
                lhs_c = vis_sorted[0][0] if len(vis_sorted) > 0 else None
                rhs_c = vis_sorted[1][0] if len(vis_sorted) > 1 else None
                
                def get_text(c_obj):
                    if not c_obj: return ""
                    curr_ids = c_obj['ids'][0]
                    if len(curr_ids) > text_cache.get(c_obj['id'], (None, 0))[1]:
                        full_text = self.model.tokenizer.decode(curr_ids, skip_special_tokens=True)
                        text_cache[c_obj['id']] = (full_text, len(curr_ids))
                        return full_text
                    return text_cache[c_obj['id']][0]

                l_raw = get_text(lhs_c).replace('\n', ' ')
                r_raw = get_text(rhs_c).replace('\n', ' ')

                # Filter Action Display based on Prompt Requirements
                # "only show branch_top1 and branch_top2 updates in the streaming view"
                if best_action_name in ['branch_top1', 'branch_top2']:
                    display_action = best_action_name
                else:
                    display_action = "" # Hide keep/prune/swap from the prominent label to reduce noise

                if ENV_MODE == 'notebook':
                    l_esc = html.escape(l_raw[-2000:])
                    r_esc = html.escape(r_raw[-2000:])
                    
                    lhs_head = f"LHS (Top-1) [ID:{lhs_c['id']}]"
                    rhs_head = f"RHS (Top-2) [ID:{rhs_c['id'] if rhs_c else '-'}]"
                    
                    css = """
                    <style>
                        .zip-dash { font-family: 'Courier New', monospace; background: #1e1e1e; color: #d4d4d4; padding: 10px; border: 1px solid #333; }
                        .zip-header { border-bottom: 2px solid #444; padding-bottom: 5px; margin-bottom: 5px; font-weight: bold; }
                        .zip-meta { color: #569cd6; border-bottom: 1px dashed #444; margin-bottom: 10px; padding-bottom: 5px;}
                        .zip-grid { display: grid; grid-template-columns: 1fr 1px 1fr; gap: 10px; }
                        .zip-col { white-space: pre-wrap; word-wrap: break-word; overflow-y: hidden; max-height: 1000px; line-height: 1.3; }
                        .zip-sep { background: #444; height: 100%; }
                    </style>
                    """
                    
                    body = f"""
                    <div class="zip-dash">
                        <div class="zip-header">
                            <div style="display:flex; justify-content:space-between;">
                                <span>{lhs_head}</span>
                                <span>{rhs_head}</span>
                            </div>
                        </div>
                        <div class="zip-meta">
                            Step {step} | Action: {display_action} | Util: {best_util:.4f} | Candidates: {len(candidates)}
                        </div>
                        <div class="zip-grid">
                            <div class="zip-col">{l_esc}</div>
                            <div class="zip-sep"></div>
                            <div class="zip-col">{r_esc}</div>
                        </div>
                    </div>
                    """
                    dashboard_widget.value = css + body

                else:
                    # CLI Renderer
                    term_w = shutil.get_terminal_size((100, 24)).columns
                    col_w = (term_w - 7) // 2 
                    
                    line_sep = "=" * term_w
                    line_head = f"{f' LHS (Top-1) [ID:{lhs_c["id"]}]':<{col_w}} || {f' RHS (Top-2) [ID:{rhs_c["id"] if rhs_c else "-"}]':<{col_w}}"
                    line_meta = f" Step {step:<4} | Action: {display_action:<11} | Util: {best_util:.4f} | Pool: {len(candidates)}"
                    line_meta = f"{line_meta:<{col_w}} || "
                    
                    tail_len = col_w * 12 
                    l_lines = textwrap.wrap(l_raw[-tail_len:], width=col_w)
                    r_lines = textwrap.wrap(r_raw[-tail_len:], width=col_w)
                    
                    max_h = max(len(l_lines), len(r_lines), 1)
                    l_lines += [" " * col_w] * (max_h - len(l_lines))
                    r_lines += [" " * col_w] * (max_h - len(r_lines))
                    
                    buffer = []
                    buffer.append(line_sep)
                    buffer.append(line_head)
                    buffer.append(line_sep)
                    buffer.append(line_meta)
                    buffer.append(line_sep)
                    for i in range(max_h):
                        buffer.append(f"{l_lines[i]:<{col_w}} || {r_lines[i]:<{col_w}}")
                    buffer.append(line_sep)
                    
                    final_str = "\n".join(buffer) + "\n"
                    
                    if last_cli_height > 0:
                        sys.stdout.write(f"\033[{last_cli_height}A") 
                        sys.stdout.write("\033[J") 
                    
                    sys.stdout.write(final_str)
                    sys.stdout.flush()
                    last_cli_height = len(buffer) + 1

        finally:
            pbar.close()

        # Final Summary
        best_traj = self.select_best_trajectory(finished_trajectories)
        score = best_traj['final_score'] if best_traj else 0.0
        
        full_text = self.model.tokenizer.decode(best_traj['ids'][0], skip_special_tokens=True)
        answer = full_text.split("assistant")[-1].strip() if "assistant" in full_text else full_text

        print("\n" + "=" * 100)
        print(f"Best Trajectory (Top-1) Confidence: {score:.2%}")
        print("=" * 100 + "\n")
        print(answer)
        print("\n" + "=" * 100)

        return finished_trajectories

    def generate(self, prompt, max_new_tokens=512, initial_samples=2):
        input_ids = self.model.tokenizer.apply_chat_template(
            [{"role": "user", "content": prompt}], tokenize=True, return_tensors="pt", add_generation_prompt=True
        ).to(self.model.device)

        candidates = []
        for i in range(initial_samples):
            candidates.append({
                'id': i, 'ids': input_ids.clone(), 'finished': False,
                'buffer': PredictionBuffer(self.config.smoothing_window),
                'joint_probs': None, 'current_len': 0
            })
        finished_trajectories = []

        for step in range(max_new_tokens):
            if not candidates: break
            active_ids = torch.cat([c['ids'] for c in candidates], dim=0)
            
            with torch.no_grad():
                outputs = self.model.base_model(active_ids)
                next_token_logits = outputs.logits[:, -1, :]
                raw_joint = self.model.get_joint_distribution(next_token_logits)

            for i, c in enumerate(candidates):
                c['buffer'].add(raw_joint[i])
                c['joint_probs'] = c['buffer'].get_smoothed()
                c['current_len'] = step
                
                valid_logits = next_token_logits[i].clone()
                valid_logits[self.model.zip_start_id : self.model.zip_start_id + self.config.total_zip_tokens] = -float('inf')
                probs = F.softmax(valid_logits, dim=-1)
                next_token = torch.multinomial(probs, 1).unsqueeze(0)
                c['ids'] = torch.cat([c['ids'], next_token], dim=1)
                if next_token.item() == self.model.tokenizer.eos_token_id:
                    c['finished'] = True
                    finished_trajectories.append(c)

            candidates = [c for c in candidates if not c['finished']]
            if not candidates: break

            # Meta-Actions (Branching/Pruning/Swapping)
            cand_metrics = []
            r_vals = (self.config.r_boundaries[:-1] + self.config.r_boundaries[1:]).to(self.model.device)/2
            for i, c in enumerate(candidates):
                qv, _ = ZIPRCMath.get_marginals(c['joint_probs'].unsqueeze(0))
                e_r = torch.sum(qv * r_vals).item()
                cand_metrics.append((i, e_r))
            
            sorted_by_reward = sorted(cand_metrics, key=lambda x: x[1], reverse=True)
            top_indices = [x[0] for x in sorted_by_reward]
            
            possible_actions = [('keep', candidates)]
            MAX_SAMPLES = 8

            if len(candidates) < MAX_SAMPLES:
                top_idx = top_indices[0]
                new_set = copy.deepcopy(candidates)
                clone = copy.deepcopy(new_set[top_idx])
                clone['id'] = max([c['id'] for c in new_set], default=0) + 1
                new_set.append(clone)
                possible_actions.append(('branch_top1', new_set))
                
                if len(candidates) >= 2 and len(candidates) + 1 < MAX_SAMPLES:
                    new_set2 = copy.deepcopy(new_set)
                    clone2 = copy.deepcopy(new_set2[top_indices[1]])
                    clone2['id'] = max([c['id'] for c in new_set2], default=0) + 1
                    new_set2.append(clone2)
                    possible_actions.append(('branch_top2', new_set2))

            if len(candidates) > 1:
                worst_idx = top_indices[-1]
                new_set = [c for i, c in enumerate(candidates) if i != worst_idx]
                possible_actions.append(('prune_bot1', new_set))

            if len(candidates) > 2:
                worst_indices = set(top_indices[-2:])
                new_set = [c for i, c in enumerate(candidates) if i not in worst_indices]
                possible_actions.append(('prune_bot2', new_set))
            
            if len(candidates) > 1 and top_indices[0] != top_indices[-1]:
                top_id = candidates[top_indices[0]]['id']
                worst_idx = top_indices[-1]
                new_set = copy.deepcopy(candidates)
                new_set = [c for i, c in enumerate(new_set) if i != worst_idx]
                source = next(c for c in new_set if c['id'] == top_id)
                clone = copy.deepcopy(source)
                clone['id'] = max([c['id'] for c in new_set], default=0) + 1
                new_set.append(clone)
                possible_actions.append(('swap', new_set))

            best_action_name, best_util, best_next_candidates = 'keep', -float('inf'), candidates
            for name, cand_set in possible_actions:
                if not cand_set: continue
                penalty = 0.0 if name == 'keep' else 0.01
                util = ZIPRCMath.compute_sampling_utility(cand_set, self.config) - penalty
                if util > best_util:
                    best_util, best_action_name, best_next_candidates = util, name, cand_set
            
            if best_action_name != 'keep':
                print(f"Step {step}: Meta-Action -> {best_action_name} (Util: {best_util:.4f}) | Pool: {len(best_next_candidates)}")
            candidates = best_next_candidates

        return finished_trajectories