import paddle
from paddlefleet.transformer import TransformerConfig
from paddlefleet.models.gpt.gpt_config import GPTConfig
from auto_configurator import (
    AutoConfigurator,
    PaddleFleetRecipe,
    generate_configs,
)

# Create model configuration
model_config = GPTConfig(
    num_hidden_layers=32,
    hidden_size=4096,
    num_attention_heads=32,
    intermediate_size=16384,
    max_sequence_length=4096,
    vocab_size=32000,
    tensor_model_parallel_size=1,
    pipeline_model_parallel_size=1,
)

# Create recipe
recipe = PaddleFleetRecipe(
    model_config=model_config,
    micro_batch_size=4,
    global_batch_size=2048,
    num_nodes=4,
    num_gpus_per_node=8,
)

# Create AutoConfigurator
runner = AutoConfigurator(
    recipe=recipe,
    path_to_logs="./logs",
    gpu_memory_gb=80,
    tensor_parallel_sizes="auto",
    pipeline_parallel_sizes="auto",
    micro_batch_sizes=[1, 2, 4, 8],
    max_steps_per_run=100,
    max_training_days=7,
    num_tokens_in_b=1400,
    vocab_size=32000,
)

# Generate configurations
base_config, configs = generate_configs(runner)
print(f"Generated {len(configs)} candidate configurations")

@dataclass
class AutoConfigurator:
    recipe: PaddleFleetRecipe           # Training recipe with model and hardware configs
    path_to_logs: str                    # Directory for saving logs

    # Hardware constraints
    gpu_memory_gb: Optional[int] = 80          # 80 GB (H100)
    tensor_parallel_sizes: Optional[List[int]] = None   # None or "auto" or explicit list
    pipeline_parallel_sizes: Optional[List[int]] = None  # None or "auto" or explicit list
    micro_batch_sizes: Optional[List[int]] = None        # None or "auto" or explicit list
    context_parallel_sizes: Optional[List[int]] = None
    expert_parallel_sizes: Optional[List[int]] = None

    # Training constraints
    num_tokens_in_b: Optional[int] = 1400    # Dataset size in billions
    tflops_per_gpu: Optional[int] = 989         # BF16 TFLOPS per GPU (H100)
    max_steps_per_run: Optional[int] = 50        # Grid search steps per config
    max_training_days: Optional[int] = 2            # Expected training days
    vocab_size: Optional[int] = 32000          # Tokenizer vocab size

    # Model calculation
    calculate_model_size: Optional[bool] = False  # Auto-calculate architecture

@dataclass
class PaddleFleetRecipe:
    model_config: object                   # TransformerConfig or subclass
    parallel_config: object | None      # Optional separate parallel config
    micro_batch_size: int = 1
    global_batch_size: int = 512
    num_nodes: int = 1
    num_gpus_per_node: int = 8
    max_steps: int | None = None
    log_dir: str | None = None

    @property
    def total_gpus(self) -> int:
        return self.num_nodes * self.num_gpus_per_node

from auto_configurator import generate_configs, AutoConfigurator, PaddleFleetRecipe

runner = AutoConfigurator(recipe=recipe, path_to_logs="./logs")
base_config, configs = generate_configs(runner)

# Iterate through generated configs
for name, config in configs.items():
    print(f"Config: {name}")
    print(f"  TP={config.tensor_parallel_size}, PP={config.pipeline_parallel_size}")

# With scoring function for top-N selection
def my_scoring_fn(cfg):
    return cfg.expert_parallel_size / cfg.tensor_parallel_size

base_config, top3 = generate_configs(runner, max_configs=3, scoring_fn=my_scoring_fn)

from auto_configurator import estimate_model_size

# Estimate model size for 7 days on 64 H100s
size = estimate_model_size(
    gpu_count=64,
    max_training_days=7,
    tflops_per_gpu=989,
    num_tokens_in_b=300,
    model_name="gpt"
)
print(f"Estimated model size: {size:.2f}B")

from auto_configurator import get_results, AutoConfigurator

runner = AutoConfigurator(...)
get_results(
    base_config=runner.recipe.model_config,
    runner_config=runner,
    path_to_save="./logs",
    output_top_n=10,
)

from paddlefleet import ModelParallelConfig
from paddlefleet.parallel_state import initialize_model_parallel
from auto_configurator import AutoConfigurator, generate_configs, PaddleFleetRecipe

# 1. Generate configurations
runner = AutoConfigurator(...)
base_config, configs = generate_configs(runner)

# 2. Select a configuration
selected_config = configs[list(configs.keys())[args.run_number - 1]]

# 3. Apply configuration to model
base_config.tensor_model_parallel_size = selected_config.tensor_parallel_size
base_config.pipeline_model_parallel_size = selected_config.pipeline_parallel_size
base_config.context_parallel_size = selected_config.context_parallel_size

# 4. Initialize parallel state
mp_config = ModelParallelConfig(
    tensor_model_parallel_size=base_config.tensor_model_parallel_size,
    pipeline_model_parallel_size=base_config.pipeline_model_parallel_size,
    context_parallel_size=base_config.context_parallel_size,
)

# 5. Start training
initialize_model_parallel(
    hcg=...  # Get your HCG from PaddlePaddle
    virtual_pipeline_model_parallel_size=base_config.virtual_pipeline_model_parallel_size,
)

# Your PaddleFleet training code here...

runner = AutoConfigurator(
    recipe=recipe,
    path_to_logs="./logs",
    # Custom search space
    tensor_parallel_sizes=[1, 2, 4, 8],
    pipeline_parallel_sizes=[1, 2, 4, 8],
    micro_batch_sizes=[2, 4, 8],
    context_parallel_sizes=[1, 2, 4],
    expert_parallel_sizes=[1, 2, 4],
)

runner = AutoConfigurator(
    recipe=recipe,
    path_to_logs="./logs",
    calculate_model_size=True,  # Enable auto-calculation
    # ... other params
)

# AutoConfigurator will infer optimal:
# - num_layers
# - hidden_size
# - num_attention_heads
# - intermediate_size