"""Command-line interface for MoltenFlow.

This module provides CLI commands for training models, generating molecules,
and optimizing molecular properties using the MoltenFlow framework.
"""

import typer

app = typer.Typer(help="MoltenFlow CLI - Property-guided molecular generation and optimization")


@app.command()
def hello():
    """Sanity check command."""
    print("MoltenFlow is installed and working.")


@app.command()
def pretrain_vae(
    config: str = typer.Option("configs/pretrain.yaml", help="Path to configuration YAML file"),
):
    """Pretrain VAE on large unlabeled SMILES corpus.

    Example:
        moltenflow pretrain-vae --config configs/pretrain.yaml
    """
    from moltenflow.training.pretrain_vae import main

    main(config_path=config)


@app.command()
def finetune_vae(
    config: str = typer.Option("configs/finetune.yaml", help="Path to configuration YAML file"),
):
    """Fine-tune pretrained VAE with property supervision.

    Example:
        moltenflow finetune-vae --config configs/finetune.yaml
    """
    from moltenflow.training.finetune_vae import main

    main(config_path=config)


@app.command()
def train_vae(
    config: str = typer.Option("configs/vae.yaml", help="Path to configuration YAML file"),
):
    """Train VAE from scratch (standard training).

    Example:
        moltenflow train-vae --config configs/vae.yaml
    """
    from moltenflow.training.train_vae import main

    main(config_path=config)


@app.command()
def train_flow(
    config: str = typer.Option("configs/latent_model.yaml", help="Path to configuration YAML file"),
):
    """Train latent flow model on VAE latents.

    Example:
        moltenflow train-flow --config configs/latent_model.yaml
    """
    from moltenflow.training.train_latent import main

    main(config_path=config)


@app.command()
def train_surrogate(
    config: str = typer.Option("configs/surrogate.yaml", help="Path to configuration YAML file"),
):
    """Train property surrogate predictor.

    Example:
        moltenflow train-surrogate --config configs/surrogate.yaml
    """
    from moltenflow.training.train_surrogate import main

    main(config_path=config)


@app.command()
def generate(
    config: str = typer.Option(
        "configs/experiments/generation_uncond.yaml", help="Path to configuration YAML file"
    ),
    guided: bool = typer.Option(False, help="Use property guidance"),
):
    """Generate molecules using VAE + flow.

    Examples:
        moltenflow generate --config configs/experiments/generation_uncond.yaml
        moltenflow generate --config configs/experiments/generation_guided.yaml --guided
    """
    if guided:
        from moltenflow.inference.generate_conditioned import main
    else:
        from moltenflow.inference.generate_uncond import main
    main(config_path=config)


@app.command()
def optimize(
    config: str = typer.Option(
        "configs/experiments/optimization_local.yaml", help="Path to configuration YAML file"
    ),
):
    """Optimize existing molecules toward target properties.

    Example:
        moltenflow optimize --config configs/experiments/optimization_local.yaml
    """
    from moltenflow.inference.optimize_conditioned import main

    main(config_path=config)


if __name__ == "__main__":
    app()
