#!/usr/bin/env python

# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Synthetic Data Generation for Hi-Robot Style Hierarchical Policy Training.

This script generates synthetic user prompts (ℓ_t) and robot utterances (u_t) for
hierarchical policy training using Qwen VLM as the generator model (pgen).

The pipeline:
1. Loads a LeRobot dataset with skill annotations (from annotate.py)
2. For each frame, generates synthetic dialogue based on:
   - Visual context (images at time t)
   - Current skill being performed
   - History of previous skills
   - High-level task description
3. Saves results as high-level tasks and updates dataset with task_index_high_level

Usage:
```bash
python examples/dataset/annotate_pgen.py \
    --repo-id lerobot/svla_so101_pickplace \
    --model Qwen/Qwen2-VL-7B-Instruct \
    --output-dir /path/to/output \
    --batch-size 1
```
"""

import argparse
import json
import re
import textwrap
from pathlib import Path
from typing import Any

import numpy as np
import pandas as pd
import torch
from PIL import Image
from rich.console import Console
from rich.progress import Progress, SpinnerColumn, TextColumn
from tqdm import tqdm

from lerobot.datasets.dataset_tools import add_features
from lerobot.datasets.lerobot_dataset import LeRobotDataset


# =============================================================================
# Prompt Template for pgen
# =============================================================================

PGEN_PROMPT_TEMPLATE = textwrap.dedent("""\
    # Role
    You are a robot-assistant dialogue generator for hierarchical robot policies.
    
    # Task
    You will receive:
    - A list of images showing the current robot scene at time t
    - The high-level task: {task_description}
    - Previous skill steps completed: {skill_history}
    - The next skill to be performed by the robot: {skill_current}
    
    # Your Goal
    Generate two things that create a natural human-robot interaction:
    1. **user_prompt**: A natural-sounding user request that logically leads to the robot 
       performing the skill "{skill_current}" given the task context and history.
    2. **robot_utterance**: A natural robot reply acknowledging or clarifying the request.
    
    # Guidelines
    - The user prompt should be grounded in the visual scene and task context
    - Vary interaction types: direct commands, implicit requests, corrections, constraints
    - Examples of user prompt styles:
      * Direct: "Can you pick up the red brick?"
      * Implicit: "I need something red for the tower"
      * Negative: "Don't pick up the blue one"
      * Constraint: "Make sure to handle it gently"
      * Correction: "Actually, move to the other box instead"
    - Robot responses should be appropriate: confirmations, clarifications, or error handling
    - Use the skill history to ensure continuity (don't repeat past actions)
    - Consider world knowledge (dietary preferences, object properties, etc.)
    
    # Scenario Types (choose one that fits):
    - **specific_object**: User specifies exact object/action
    - **negative_task**: User says what NOT to do
    - **situated_correction**: User adjusts based on current state
    - **implicit_request**: User implies need without direct command
    - **constraint_based**: User adds specific constraints
    
    # Response Types (choose one that fits):
    - **confirmation**: Simple "OK, I'll do X"
    - **clarification**: "Just to confirm, you want me to..."
    - **acknowledgment**: "Got it, [doing action]"
    - **constraint_acknowledgment**: "Sure, I'll [action] while [constraint]"
    
    # Output Format
    Respond ONLY with valid JSON:
    {{
      "scenario_type": "one of the types above",
      "response_type": "one of the types above", 
      "user_prompt": "natural user request here",
      "robot_utterance": "natural robot response here"
    }}
    
    The responses must be grounded in the visual scene, the task, and the skill history.
    Make it sound like a real human-robot interaction.
    """)


def construct_prompt(
    task_description: str,
    skill_history: list[str],
    skill_current: str,
) -> str:
    """
    Construct the text prompt for pgen.
    
    Args:
        task_description: High-level task description
        skill_history: List of previously completed skills
        skill_current: Current skill to be performed
        
    Returns:
        Formatted prompt string
    """
    # Format skill history nicely
    if skill_history:
        history_str = ", ".join(f'"{s}"' for s in skill_history[-5:])  # Last 5 for context
        if len(skill_history) > 5:
            history_str = f"... {history_str}"
    else:
        history_str = "None (starting the task)"
    
    return PGEN_PROMPT_TEMPLATE.format(
        task_description=task_description,
        skill_history=history_str,
        skill_current=skill_current,
    )


# =============================================================================
# Qwen VLM Interface
# =============================================================================

class QwenPgen:
    """Qwen VLM wrapper for synthetic dialogue generation."""
    
    def __init__(
        self,
        model_name: str,
        device: str = "cuda",
        torch_dtype: torch.dtype = torch.bfloat16,
        temperature: float = 0.7,
    ):
        from qwen_vl_utils import process_vision_info
        from transformers import AutoProcessor, Qwen2VLForConditionalGeneration
        
        self.console = Console()
        self.device = device
        self.model_name = model_name
        self.temperature = temperature
        self.process_vision_info = process_vision_info
        
        self.console.print(f"[cyan]Loading Qwen model: {model_name}...[/cyan]")
        
        # Load model based on name
        if "qwen3" in model_name.lower():
            from transformers import Qwen3VLMoeForConditionalGeneration
            self.model = Qwen3VLMoeForConditionalGeneration.from_pretrained(
                model_name, torch_dtype=torch_dtype, device_map=device, trust_remote_code=True
            )
        else:
            self.model = Qwen2VLForConditionalGeneration.from_pretrained(
                model_name, torch_dtype=torch_dtype, device_map=device, trust_remote_code=True
            )
        
        self.processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True)
        self.console.print(f"[green]✓ Model loaded successfully on {device}[/green]")
    
    def call_qwen(
        self,
        images: list[Image.Image | str],
        prompt: str,
    ) -> dict[str, str]:
        """
        Call Qwen VLM to generate synthetic dialogue.
        
        Args:
            images: List of PIL Images or image paths
            prompt: Text prompt for generation
            
        Returns:
            Dictionary with keys: scenario_type, response_type, user_prompt, robot_utterance
        """
        # Build messages with images and text
        content = []
        for img in images:
            if isinstance(img, str):
                content.append({"type": "image", "image": img})
            else:
                # PIL Image - need to save temporarily or convert
                content.append({"type": "image", "image": img})
        
        content.append({"type": "text", "text": prompt})
        
        messages = [
            {
                "role": "user",
                "content": content,
            }
        ]
        
        # Process inputs
        text = self.processor.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )
        image_inputs, video_inputs = self.process_vision_info(messages)
        
        inputs = self.processor(
            text=[text],
            images=image_inputs,
            videos=video_inputs,
            padding=True,
            return_tensors="pt",
        ).to(self.device)
        
        # Generate
        with torch.no_grad():
            generated_ids = self.model.generate(
                **inputs,
                max_new_tokens=512,
                do_sample=True,
                temperature=self.temperature,
            )
        
        # Decode response
        response = self.processor.batch_decode(
            [out[len(inp):] for inp, out in zip(inputs.input_ids, generated_ids)],
            skip_special_tokens=True,
        )[0].strip()
        
        return self._parse_response(response)
    
    def _parse_response(self, response: str) -> dict[str, str]:
        """Parse JSON response from model."""
        # Extract JSON from response
        if "```json" in response:
            response = response.split("```json")[1].split("```")[0]
        elif "```" in response:
            response = response.split("```")[1].split("```")[0]
        
        try:
            data = json.loads(response)
            return {
                "scenario_type": data.get("scenario_type", "specific_object"),
                "response_type": data.get("response_type", "confirmation"),
                "user_prompt": data.get("user_prompt", ""),
                "robot_utterance": data.get("robot_utterance", ""),
            }
        except json.JSONDecodeError:
            # Try to find JSON object in response
            match = re.search(r"\{.*\}", response, re.DOTALL)
            if match:
                data = json.loads(match.group())
                return {
                    "scenario_type": data.get("scenario_type", "specific_object"),
                    "response_type": data.get("response_type", "confirmation"),
                    "user_prompt": data.get("user_prompt", ""),
                    "robot_utterance": data.get("robot_utterance", ""),
                }
            
            raise ValueError(f"Could not parse response: {response[:200]}...")


# =============================================================================
# Annotation Pipeline
# =============================================================================

def load_skills_metadata(dataset_root: Path) -> dict | None:
    """Load skills.json metadata from annotated dataset."""
    skills_path = dataset_root / "meta" / "skills.json"
    if skills_path.exists():
        with open(skills_path) as f:
            return json.load(f)
    return None


def get_skill_at_timestamp(skills: list[dict], timestamp: float) -> str | None:
    """Find which skill covers a given timestamp."""
    for skill in skills:
        if skill["start"] <= timestamp < skill["end"]:
            return skill["name"]
        # Handle last frame
        if timestamp >= skill["end"] and skill == skills[-1]:
            return skill["name"]
    return skills[-1]["name"] if skills else None


def annotate_sample(
    pgen: QwenPgen,
    images: list[Image.Image | str],
    task_description: str,
    skill_history: list[str],
    skill_current: str,
) -> dict[str, str]:
    """
    Generate synthetic dialogue for a single sample.
    
    Args:
        pgen: Qwen model wrapper
        images: List of images at current timestep
        task_description: High-level task description
        skill_history: Previous skills completed
        skill_current: Current skill being performed
        
    Returns:
        Dictionary with generated dialogue
    """
    prompt = construct_prompt(task_description, skill_history, skill_current)
    result = pgen.call_qwen(images, prompt)
    return result


def generate_synthetic_data(
    dataset: LeRobotDataset,
    pgen: QwenPgen,
    skills_metadata: dict,
    image_keys: list[str],
    sample_interval_seconds: float = 1.0,
    console: Console | None = None,
) -> tuple[pd.DataFrame, np.ndarray, list[dict]]:
    """
    Generate synthetic dialogue data for entire dataset.
    
    This function processes ALL frames in the dataset, but only calls the VLM
    at specified intervals (sample_interval_seconds). Frames between samples
    inherit the task_index from the most recent sample.
    
    Args:
        dataset: LeRobot dataset with skill annotations
        pgen: Qwen model wrapper
        skills_metadata: Loaded skills.json metadata
        image_keys: List of image observation keys to use
        sample_interval_seconds: Generate dialogue every N seconds (default: 1.0)
        console: Rich console for logging
        
    Returns:
        Tuple of (tasks_df, task_indices_array, debug_outputs)
        - tasks_df: DataFrame with high-level tasks (user_prompt, robot_utterance, etc.)
        - task_indices_array: Array of task indices for each frame (full dataset length)
        - debug_outputs: List of debug dictionaries (only for sampled frames)
    """
    if console is None:
        console = Console()
    
    # Extract metadata
    coarse_description = skills_metadata.get("coarse_description", "Complete the task")
    episodes = skills_metadata.get("episodes", {})
    
    # Track unique high-level tasks
    high_level_tasks = {}  # (user_prompt, robot_utterance, skill) -> task_index
    task_index_counter = 0  # Start at 0
    
    # Array to store task index for each frame - MUST match full dataset length
    full_dataset_length = len(dataset)
    task_indices = np.zeros(full_dataset_length, dtype=np.int64)
    
    # For debugging - save to JSONL
    debug_outputs = []
    
    # Track sampling
    last_sample_timestamp = {}  # episode_idx -> last sampled timestamp
    last_task_index = {}  # episode_idx -> last generated task_index
    frames_sampled = 0
    
    console.print(f"[cyan]Processing all {full_dataset_length} frames from {dataset.meta.total_episodes} episodes...[/cyan]")
    console.print(f"[cyan]Sampling interval: {sample_interval_seconds}s (fps: {dataset.meta.fps})[/cyan]")
    
    # Process each frame in the FULL dataset
    for frame_idx in tqdm(range(full_dataset_length), desc="Generating synthetic dialogue"):
        try:
            # Get frame data
            frame = dataset[frame_idx]
            episode_idx = frame["episode_index"].item()
            timestamp = frame["timestamp"].item()
            
            # Get episode skills
            episode_key = str(episode_idx)
            if episode_key not in episodes:
                console.print(f"[yellow]Warning: Episode {episode_idx} not in skills metadata[/yellow]")
                continue
            
            episode_data = episodes[episode_key]
            skills = episode_data.get("skills", [])
            description = episode_data.get("description", coarse_description)
            
            # Find current skill
            current_skill = get_skill_at_timestamp(skills, timestamp)
            if current_skill is None:
                console.print(f"[yellow]Warning: No skill found for timestamp {timestamp}[/yellow]")
                continue
            
            # Determine if we should sample this frame
            should_sample = False
            
            # Always sample first frame of an episode
            if episode_idx not in last_sample_timestamp:
                should_sample = True
                last_sample_timestamp[episode_idx] = timestamp
            else:
                # Sample if enough time has passed
                time_since_last = timestamp - last_sample_timestamp[episode_idx]
                if time_since_last >= sample_interval_seconds:
                    should_sample = True
                    last_sample_timestamp[episode_idx] = timestamp
            
            # If not sampling, reuse last task index for this episode
            if not should_sample:
                if episode_idx in last_task_index:
                    task_indices[frame_idx] = last_task_index[episode_idx]
                continue
            
            # Sample this frame - generate synthetic dialogue
            frames_sampled += 1
            
            # Build skill history (all skills before current timestamp)
            skill_history = []
            for skill in skills:
                if skill["end"] <= timestamp:
                    skill_history.append(skill["name"])
            
            # Load images
            images = []
            for img_key in image_keys:
                if img_key in frame:
                    # Frame images are tensors (C, H, W) in [0, 1]
                    img_tensor = frame[img_key]
                    if len(img_tensor.shape) == 4:  # (T, C, H, W)
                        img_tensor = img_tensor[-1]  # Take last frame
                    
                    # Convert to PIL Image
                    img_array = (img_tensor.permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8)
                    img_pil = Image.fromarray(img_array)
                    images.append(img_pil)
            
            if not images:
                console.print(f"[yellow]Warning: No images found for frame {frame_idx}[/yellow]")
                continue
            
            # Generate synthetic dialogue
            result = annotate_sample(
                pgen=pgen,
                images=images,
                task_description=description,
                skill_history=skill_history,
                skill_current=current_skill,
            )
            
            # Create unique task key
            task_key = (
                result["user_prompt"],
                result["robot_utterance"],
                current_skill,
                result["scenario_type"],
                result["response_type"],
            )
            
            # Assign or create task index
            if task_key not in high_level_tasks:
                high_level_tasks[task_key] = task_index_counter
                task_index_counter += 1
            
            current_task_idx = high_level_tasks[task_key]
            task_indices[frame_idx] = current_task_idx
            last_task_index[episode_idx] = current_task_idx
            
            # Save for debugging
            debug_outputs.append({
                "episode_id": int(episode_idx),
                "frame_index": frame_idx,
                "timestamp": float(timestamp),
                "skill_current": current_skill,
                "skill_history": skill_history,
                "task_description": description,
                "sampled": True,
                **result,
            })
            
        except Exception as e:
            console.print(f"[red]Error processing frame {frame_idx}: {e}[/red]")
            continue
    
    console.print(f"[green]✓ Sampled {frames_sampled} frames out of {full_dataset_length} total ({frames_sampled/full_dataset_length*100:.1f}%)[/green]")
    
    # Create tasks DataFrame
    tasks_data = []
    for task_key, task_idx in sorted(high_level_tasks.items(), key=lambda x: x[1]):
        user_prompt, robot_utterance, skill, scenario_type, response_type = task_key
        tasks_data.append({
            "task": f"{user_prompt} | {robot_utterance}",
            "task_index": task_idx,
            "user_prompt": user_prompt,
            "robot_utterance": robot_utterance,
            "skill": skill,
            "scenario_type": scenario_type,
            "response_type": response_type,
        })
    
    tasks_df = pd.DataFrame(tasks_data).set_index("task")
    
    console.print(f"[green]✓ Generated {len(high_level_tasks)} unique high-level tasks[/green]")
    
    return tasks_df, task_indices, debug_outputs


def save_high_level_tasks(
    tasks_df: pd.DataFrame,
    dataset_root: Path,
    console: Console | None = None,
) -> None:
    """Save high-level tasks to tasks_high_level.parquet."""
    if console is None:
        console = Console()
    
    output_path = dataset_root / "meta" / "tasks_high_level.parquet"
    output_path.parent.mkdir(parents=True, exist_ok=True)
    
    tasks_df.to_parquet(output_path, engine="pyarrow", compression="snappy")
    console.print(f"[green]✓ Saved high-level tasks to {output_path}[/green]")


def save_debug_outputs(
    debug_outputs: list[dict],
    dataset_root: Path,
    console: Console | None = None,
) -> None:
    """Save debug outputs to JSONL file."""
    if console is None:
        console = Console()
    
    output_path = dataset_root / "meta" / "syn_annotations.jsonl"
    output_path.parent.mkdir(parents=True, exist_ok=True)
    
    with open(output_path, "w") as f:
        for item in debug_outputs:
            f.write(json.dumps(item) + "\n")
    
    console.print(f"[green]✓ Saved debug annotations to {output_path}[/green]")


# =============================================================================
# Main Entry Point
# =============================================================================

def main():
    """Main entry point for synthetic data generation."""
    parser = argparse.ArgumentParser(
        description="Generate synthetic dialogue data for hierarchical robot policies",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=textwrap.dedent("""\
            Examples:
              # Generate synthetic data for a dataset
              python annotate_pgen.py --repo-id lerobot/svla_so101_pickplace \\
                  --model Qwen/Qwen2-VL-7B-Instruct \\
                  --output-dir ./output
              
              # Use Qwen3 model with custom parameters
              python annotate_pgen.py --repo-id lerobot/svla_so101_pickplace \\
                  --model Qwen/Qwen3-VL-30B-A3B-Instruct \\
                  --temperature 0.8 \\
                  --batch-size 1
        """),
    )
    
    # Data source
    data_group = parser.add_mutually_exclusive_group(required=True)
    data_group.add_argument("--data-dir", type=str, help="Path to local LeRobot dataset")
    data_group.add_argument("--repo-id", type=str, help="HuggingFace Hub dataset repository ID")
    
    # Model configuration
    parser.add_argument(
        "--model",
        type=str,
        default="Qwen/Qwen2-VL-7B-Instruct",
        help="VLM model to use (default: Qwen/Qwen2-VL-7B-Instruct)",
    )
    parser.add_argument(
        "--device",
        type=str,
        default="cuda",
        help="Device to run model on (default: cuda)",
    )
    parser.add_argument(
        "--dtype",
        type=str,
        default="bfloat16",
        choices=["bfloat16", "float16", "float32"],
        help="Model dtype (default: bfloat16)",
    )
    parser.add_argument(
        "--temperature",
        type=float,
        default=0.7,
        help="Sampling temperature (default: 0.7)",
    )
    
    # Processing options
    parser.add_argument(
        "--batch-size",
        type=int,
        default=1,
        help="Batch size for processing (default: 1) [currently unused]",
    )
    parser.add_argument(
        "--num-image-views-per-sample",
        type=int,
        default=1,
        help="Number of camera views to use per sample (default: 1)",
    )
    parser.add_argument(
        "--sample-interval",
        type=float,
        default=1.0,
        help="Generate dialogue every N seconds (default: 1.0). Frames between samples reuse the last generated dialogue. "
             "Use larger intervals (e.g., 2.0 or 5.0) for faster processing during testing.",
    )
    
    # Output options
    parser.add_argument(
        "--output-dir",
        type=str,
        default=None,
        help="Output directory for modified dataset",
    )
    parser.add_argument(
        "--push-to-hub",
        action="store_true",
        help="Push modified dataset to HuggingFace Hub",
    )
    
    args = parser.parse_args()
    console = Console()
    
    # Load dataset
    console.print("[cyan]Loading dataset...[/cyan]")
    if args.data_dir:
        dataset = LeRobotDataset(repo_id="local/dataset", root=args.data_dir)
        dataset_root = Path(args.data_dir)
    else:
        dataset = LeRobotDataset(repo_id=args.repo_id)
        dataset_root = dataset.root
    
    console.print(f"[green]✓ Loaded dataset with {len(dataset)} frames[/green]")
    
    # Load skills metadata
    console.print("[cyan]Loading skills metadata...[/cyan]")
    skills_metadata = load_skills_metadata(dataset_root)
    if skills_metadata is None:
        console.print("[red]Error: No skills.json found. Run annotate.py first![/red]")
        return
    
    console.print(f"[green]✓ Loaded skills for {len(skills_metadata.get('episodes', {}))} episodes[/green]")
    
    # Initialize model
    dtype_map = {
        "bfloat16": torch.bfloat16,
        "float16": torch.float16,
        "float32": torch.float32,
    }
    torch_dtype = dtype_map[args.dtype]
    
    console.print(f"[cyan]Initializing {args.model}...[/cyan]")
    pgen = QwenPgen(
        model_name=args.model,
        device=args.device,
        torch_dtype=torch_dtype,
        temperature=args.temperature,
    )
    
    # Get image keys
    image_keys = dataset.meta.camera_keys[:args.num_image_views_per_sample]
    console.print(f"[cyan]Using image keys: {image_keys}[/cyan]")
    
    # Generate synthetic data
    tasks_df, task_indices, debug_outputs = generate_synthetic_data(
        dataset=dataset,
        pgen=pgen,
        skills_metadata=skills_metadata,
        image_keys=image_keys,
        sample_interval_seconds=args.sample_interval,
        console=console,
    )
    
    # Save high-level tasks
    save_high_level_tasks(tasks_df, dataset_root, console)
    save_debug_outputs(debug_outputs, dataset_root, console)
    
    # Add task_index_high_level feature to dataset
    console.print("[cyan]Adding task_index_high_level feature to dataset...[/cyan]")
    
    # Determine output directory
    if args.output_dir:
        output_dir = Path(args.output_dir)
        repo_id = f"{dataset.repo_id}_with_high_level_tasks"
    else:
        output_dir = None
        repo_id = f"{dataset.repo_id}_with_high_level_tasks"
    
    # Add feature using dataset_tools
    feature_info = {
        "dtype": "int64",
        "shape": (1,),
        "names": None,
    }
    breakpoint()
    new_dataset = add_features(
        dataset=dataset,
        features={
            "task_index_high_level": (task_indices, feature_info),
        },
        output_dir=output_dir,
        repo_id=repo_id,
    )
    
    console.print(f"[bold green]✓ Successfully added task_index_high_level feature![/bold green]")
    console.print(f"  New dataset saved to: {new_dataset.root}")
    console.print(f"  Total high-level tasks: {len(tasks_df)}")
    
    # Push to hub if requested
    if args.push_to_hub:
        if args.data_dir:
            console.print("[yellow]Warning: --push-to-hub requires --repo-id, skipping...[/yellow]")
        else:
            console.print("[cyan]Pushing to HuggingFace Hub...[/cyan]")
            try:
                new_dataset.push_to_hub(push_videos=False)
                console.print(f"[green]✓ Pushed to {repo_id}[/green]")
            except Exception as e:
                console.print(f"[red]Push failed: {e}[/red]")


if __name__ == "__main__":
    main()