"""
Neon VLA Training Script.

Trains the action heads (and optionally LoRA on backbone) using mixed data soup.
Designed to run on HuggingFace GPU Jobs or local CUDA machines.

Usage:
    # HuggingFace GPU Job (A100)
    hf jobs uv run --flavor a100-large --secrets HF_TOKEN --timeout 8h scripts/train_neon.py

    # Local
    python scripts/train_neon.py --config configs/g1_arms_only.yaml

    # As module
    from neon.training.train import train
    train(config)
"""

from __future__ import annotations

import json
import logging
import os
import sys
import time
from pathlib import Path
from typing import Dict, Optional

import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader

from neon.data.action_space import G1ActionSpace
from neon.data.data_soup import DataSoupConfig, DataSoupDataset
from neon.model.neon_vla import NeonConfig, NeonVLA
from neon.training.config import TrainConfig

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger(__name__)


class NeonTrainer:
    """Training loop for the Neon VLA model.

    Handles:
    - Data loading with data soup mixing
    - Action head training (always)
    - Optional LoRA fine-tuning on backbone
    - Gradient checkpointing for memory efficiency
    - Checkpointing and hub upload
    """

    def __init__(self, config: TrainConfig):
        self.config = config
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        if not torch.cuda.is_available():
            logger.warning("No CUDA device found! Training will be very slow.")

        # Build model
        logger.info("Building NeonVLA model...")
        self.model = NeonVLA(config.model)

        # Load backbone
        logger.info("Loading video backbone...")
        self.model.load_backbone()

        # Move action heads to device
        self.model.action_head.to(self.device)
        self.model.fusion.to(self.device)
        if self.model.proprio_encoder is not None:
            self.model.proprio_encoder.to(self.device)

        # Apply LoRA to backbone if requested
        if config.use_lora and not config.model.backbone.freeze_backbone:
            self._apply_lora()

        # Build optimizer (only for trainable parameters)
        trainable_params = [p for p in self.model.parameters() if p.requires_grad]
        self.optimizer = self._build_optimizer(trainable_params)
        self.scheduler = self._build_scheduler()

        # Print parameter counts
        total = sum(p.numel() for p in self.model.parameters())
        trainable = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
        logger.info(f"Parameters: {trainable:,} trainable / {total:,} total ({100*trainable/total:.2f}%)")

    def _apply_lora(self):
        """Apply LoRA to the backbone for fine-tuning."""
        try:
            from peft import LoraConfig, get_peft_model

            lora_config = LoraConfig(
                r=self.config.lora_r,
                lora_alpha=self.config.lora_alpha,
                target_modules=self.config.lora_target_modules,
                lora_dropout=self.config.lora_dropout,
                bias="none",
                task_type="FEATURE_EXTRACTION",
            )
            self.model.backbone._model = get_peft_model(self.model.backbone._model, lora_config)
            logger.info(f"LoRA applied: r={self.config.lora_r}, alpha={self.config.lora_alpha}")
        except ImportError:
            logger.warning("peft not installed, skipping LoRA")

    def _build_optimizer(self, params):
        """Build optimizer."""
        if self.config.optim == "paged_adamw_8bit":
            try:
                import bitsandbytes as bnb
                return bnb.optim.PagedAdamW8bit(
                    params, lr=self.config.learning_rate, weight_decay=self.config.weight_decay
                )
            except ImportError:
                logger.warning("bitsandbytes not available, using AdamW")

        return torch.optim.AdamW(
            params, lr=self.config.learning_rate, weight_decay=self.config.weight_decay
        )

    def _build_scheduler(self):
        """Build learning rate scheduler."""
        from torch.optim.lr_scheduler import CosineAnnealingLR

        # Estimate total steps
        total_steps = 10000  # Will be updated once we know dataset size
        return CosineAnnealingLR(self.optimizer, T_max=total_steps, eta_min=1e-6)

    def build_dataloader(self) -> DataLoader:
        """Build the training data loader."""
        dataset = DataSoupDataset(self.config.data, self.model.action_space)
        sampler = dataset.get_sampler()

        return DataLoader(
            dataset,
            batch_size=self.config.batch_size,
            sampler=sampler,
            num_workers=self.config.num_workers,
            pin_memory=True,
            collate_fn=self._collate_fn,
        )

    def _collate_fn(self, batch):
        """Custom collate for mixed data types.

        Handles three source types:
        - Standard (lerobot, custom): images + actions + language → full VLA loss
        - voice_commands (language_only=True): language only → language conditioning loss
        - stereo4d (visual_only=True): images only → visual backbone loss
        """
        collated = {
            "images": [item["image"] for item in batch],
            "video_frames": [item["video_frames"] for item in batch],
            "languages": [item["language"] for item in batch],
            "target_actions": torch.from_numpy(
                np.stack([item["target_actions"] for item in batch])
            ).float().to(self.device),
            "action_masks": torch.from_numpy(
                np.stack([item["action_mask"] for item in batch])
            ).float().to(self.device),
            # Source type flags for per-sample loss masking
            "language_only": [item.get("language_only", False) for item in batch],
            "visual_only": [item.get("visual_only", False) for item in batch],
        }

        # Proprioception (may be None for some sources)
        proprios = [item["proprioception"] for item in batch]
        if any(p is not None for p in proprios):
            # Replace None with zeros
            dim = self.model.action_space.action_dim
            proprios = [p if p is not None else np.zeros(dim, dtype=np.float32) for p in proprios]
            collated["proprioception"] = torch.from_numpy(np.stack(proprios)).float().to(self.device)
        else:
            collated["proprioception"] = None

        return collated

    def train_step(self, batch: Dict) -> Dict[str, float]:
        """Single training step.

        Handles three source types with different loss computation:
        - Standard: full VLA loss (visual encoding + action prediction)
        - language_only: skip — no images or actions to learn from,
          but language is already processed by backbone during encoding
        - visual_only: visual encoding only, zero action loss (mask out actions)

        Returns dict with loss values.
        """
        self.model.train()

        # Forward pass
        # For now, process images one-by-one (batch processing for VLMs is complex)
        total_loss = torch.tensor(0.0, device=self.device)
        batch_size = len(batch["images"])
        action_samples = 0  # Count samples that contribute action loss

        for i in range(batch_size):
            is_language_only = batch["language_only"][i] if "language_only" in batch else False
            is_visual_only = batch["visual_only"][i] if "visual_only" in batch else False

            # Skip language-only samples (no images, no actions to learn from)
            # Their language distribution is learned implicitly when paired with
            # visual sources in the same batch via the language encoder
            if is_language_only:
                continue

            image = batch["images"][i]
            video = batch["video_frames"][i]
            text = batch["languages"][i]
            proprio = batch["proprioception"][i:i+1] if batch["proprioception"] is not None else None
            target = batch["target_actions"][i:i+1]
            mask = batch["action_masks"][i:i+1]

            # For visual-only samples (stereo4d): zero out action mask
            # This means the backbone sees the images but no action loss is backpropagated.
            # The backbone still learns visual features via the forward pass gradients
            # flowing through the frozen backbone (if unfrozen) or through LoRA.
            if is_visual_only:
                mask = torch.zeros_like(mask)

            losses = self.model.compute_loss(
                images=[image] if image else None,
                video_frames=video if video else None,
                text=text,
                proprioception=proprio,
                target_actions=target,
                action_mask=mask,
            )
            total_loss = total_loss + losses["loss"]
            action_samples += 1

        # Avoid division by zero if entire batch is language-only
        if action_samples > 0:
            total_loss = total_loss / action_samples
        else:
            # Pure language batch — return zero loss (no backward)
            return {"loss": 0.0, "skipped": True}

        # Backward
        total_loss.backward()

        # Gradient clipping
        if self.config.max_grad_norm > 0:
            torch.nn.utils.clip_grad_norm_(
                [p for p in self.model.parameters() if p.requires_grad],
                self.config.max_grad_norm,
            )

        # Optimizer step
        self.optimizer.step()
        self.scheduler.step()
        self.optimizer.zero_grad()

        return {"loss": total_loss.item()}

    def train(self, dataloader: Optional[DataLoader] = None):
        """Full training loop."""
        if dataloader is None:
            dataloader = self.build_dataloader()

        total_steps = len(dataloader) * self.config.epochs // self.config.gradient_accumulation_steps
        logger.info(f"Starting training: {total_steps} steps, {self.config.epochs} epochs")
        logger.info(f"  Batch size: {self.config.batch_size} x {self.config.gradient_accumulation_steps} accum")

        output_dir = Path(self.config.output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)

        global_step = 0
        best_loss = float("inf")
        loss_history = []

        for epoch in range(self.config.epochs):
            epoch_loss = 0.0
            epoch_steps = 0

            for batch_idx, batch in enumerate(dataloader):
                metrics = self.train_step(batch)
                epoch_loss += metrics["loss"]
                epoch_steps += 1
                global_step += 1

                # Logging
                if global_step % self.config.logging_steps == 0:
                    avg_loss = epoch_loss / epoch_steps
                    lr = self.scheduler.get_last_lr()[0]
                    logger.info(
                        f"[Epoch {epoch+1}/{self.config.epochs}] "
                        f"Step {global_step}/{total_steps} "
                        f"Loss: {metrics['loss']:.4f} (avg: {avg_loss:.4f}) "
                        f"LR: {lr:.2e}"
                    )
                    loss_history.append({"step": global_step, "loss": metrics["loss"]})

                # Checkpointing
                if global_step % self.config.save_steps == 0:
                    ckpt_dir = output_dir / f"checkpoint-{global_step}"
                    self.model.save_pretrained(str(ckpt_dir))
                    logger.info(f"Saved checkpoint: {ckpt_dir}")

            # End of epoch
            avg_loss = epoch_loss / max(epoch_steps, 1)
            logger.info(f"Epoch {epoch+1} complete. Average loss: {avg_loss:.4f}")

            if avg_loss < best_loss:
                best_loss = avg_loss
                self.model.save_pretrained(str(output_dir / "best"))
                logger.info(f"New best model saved (loss: {best_loss:.4f})")

        # Final save
        self.model.save_pretrained(str(output_dir / "final"))
        logger.info(f"Training complete! Best loss: {best_loss:.4f}")

        # Save training stats
        stats = {
            "best_loss": best_loss,
            "total_steps": global_step,
            "epochs": self.config.epochs,
            "loss_history": loss_history,
        }
        with open(output_dir / "training_stats.json", "w") as f:
            json.dump(stats, f, indent=2)

        # Push to hub
        if self.config.push_to_hub:
            self._push_to_hub(output_dir)

        return stats

    def _push_to_hub(self, output_dir: Path):
        """Push trained model to HuggingFace Hub."""
        try:
            from huggingface_hub import HfApi, login

            token = os.getenv("HF_TOKEN")
            if token:
                login(token=token)

            api = HfApi()
            api.upload_folder(
                folder_path=str(output_dir / "best"),
                repo_id=self.config.hub_model_id,
                repo_type="model",
                token=token,
            )
            logger.info(f"Pushed to Hub: https://huggingface.co/{self.config.hub_model_id}")
        except Exception as e:
            logger.error(f"Failed to push to Hub: {e}")


def train(config: Optional[TrainConfig] = None):
    """Entry point for training."""
    if config is None:
        from neon.training.config import default_arms_only_config
        config = default_arms_only_config()

    trainer = NeonTrainer(config)
    return trainer.train()


def main():
    """CLI entry point."""
    import argparse

    parser = argparse.ArgumentParser(description="Train Neon VLA")
    parser.add_argument("--config", type=str, help="Path to YAML config file")
    parser.add_argument("--backbone", type=str, default="Qwen/Qwen2.5-VL-3B-Instruct")
    parser.add_argument("--mode", type=str, default="arms_only", choices=["arms_only", "upper_body", "whole_body"])
    parser.add_argument("--epochs", type=int, default=3)
    parser.add_argument("--batch-size", type=int, default=4)
    parser.add_argument("--output", type=str, default="cagataydev/neon-g1-v1")
    args = parser.parse_args()

    if args.config:
        # Load from YAML
        import yaml
        with open(args.config) as f:
            config_dict = yaml.safe_load(f)
        # TODO: deserialize TrainConfig from dict
        config = TrainConfig()
    else:
        from neon.model.video_backbone import BackboneConfig
        config = TrainConfig(
            model=NeonConfig(
                backbone=BackboneConfig(model_id=args.backbone, load_in_4bit=True, freeze_backbone=True),
                control_mode=args.mode,
            ),
            epochs=args.epochs,
            batch_size=args.batch_size,
            hub_model_id=args.output,
        )

    train(config)


if __name__ == "__main__":
    main()
