"""Standalone GPU training script for RunPod.

Trains sentence-transformer embeddings and a Neural Collaborative Filtering
(NCF) model on exported movie/rating data.  Fully self-contained -- does NOT
import any app modules so it can run on a bare RunPod instance with only
PyTorch and sentence-transformers installed.

Usage:
    python train_gpu.py
    python train_gpu.py --data-dir /data/export --output-dir /data/models
    python train_gpu.py --device cuda:0 --epochs 80 --batch-size 512
"""

from __future__ import annotations

import argparse
import json
import time
from pathlib import Path
from typing import Any

import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
from sentence_transformers import SentenceTransformer


# ---------------------------------------------------------------------------
# Document builder -- mirrors feature_engineering.build_movie_document
# ---------------------------------------------------------------------------

def build_movie_document(movie: dict[str, Any]) -> str:
    """Build a natural-language document from movie metadata.

    Uses natural language format optimized for sentence transformers,
    matching ``app.ml.embedding_engine.build_movie_document``.
    """
    parts: list[str] = []

    if movie.get("overview"):
        parts.append(movie["overview"])

    if movie.get("genres"):
        genres = movie["genres"] if isinstance(movie["genres"], list) else []
        if genres:
            parts.append(f"Genres: {', '.join(str(g) for g in genres)}.")

    if movie.get("director"):
        parts.append(f"Director: {movie['director']}.")

    if movie.get("cast"):
        cast = movie["cast"] if isinstance(movie["cast"], list) else []
        cast_names: list[str] = []
        for c in cast[:5]:
            if isinstance(c, dict):
                name = c.get("name")
                if name:
                    cast_names.append(str(name))
            elif isinstance(c, str):
                cast_names.append(c)
        if cast_names:
            parts.append(f"Cast: {', '.join(cast_names)}.")

    return " ".join(parts)


# ---------------------------------------------------------------------------
# Embedding training
# ---------------------------------------------------------------------------

def train_embeddings(
    movies: list[dict[str, Any]],
    output_dir: Path,
    device: str,
) -> None:
    """Encode all movies with a sentence-transformer and persist the result.

    Saves:
        - ``embeddings.npy``       -- (N, D) float32 array
        - ``embeddings_metadata.json`` -- movie id list + model info
    """
    print(f"\n{'='*60}")
    print("Training embeddings")
    print(f"{'='*60}")
    t0 = time.time()

    print(f"Loading SentenceTransformer('all-MiniLM-L6-v2') on {device} ...")
    model = SentenceTransformer("all-MiniLM-L6-v2", device=device)

    print(f"Building documents for {len(movies)} movies ...")
    documents = [build_movie_document(m) for m in movies]

    empty_count = sum(1 for d in documents if not d.strip())
    if empty_count:
        print(f"  Warning: {empty_count} movies produced empty documents")

    print("Encoding ...")
    embeddings = model.encode(
        documents,
        batch_size=64,
        show_progress_bar=True,
        convert_to_numpy=True,
    )

    embeddings_path = output_dir / "embeddings.npy"
    np.save(embeddings_path, embeddings)
    print(f"Saved embeddings ({embeddings.shape}) to {embeddings_path}")

    metadata = {
        "model_name": "all-MiniLM-L6-v2",
        "embedding_dim": int(embeddings.shape[1]),
        "n_movies": len(movies),
        "movie_ids": [m["id"] for m in movies],
    }
    metadata_path = output_dir / "embeddings_metadata.json"
    with open(metadata_path, "w") as f:
        json.dump(metadata, f, indent=2)
    print(f"Saved metadata to {metadata_path}")

    elapsed = time.time() - t0
    print(f"Embeddings done in {elapsed:.1f}s")


# ---------------------------------------------------------------------------
# NCF model definition
# ---------------------------------------------------------------------------

class NCFModel(nn.Module):
    """Neural Collaborative Filtering (GMF + MLP).

    Architecture mirrors the project's collaborative model design:
      - GMF path:  element-wise product of user/item embeddings
      - MLP path:  concatenated embeddings -> [128, 64, 32] with ReLU + Dropout
      - Output:    sigmoid(combined) * 9 + 1  (maps to 1-10 rating scale)
    """

    def __init__(
        self,
        n_users: int,
        n_items: int,
        embed_dim: int = 32,
        mlp_layers: tuple[int, ...] = (128, 64, 32),
        dropout: float = 0.2,
    ):
        super().__init__()

        # GMF embeddings
        self.gmf_user_embed = nn.Embedding(n_users, embed_dim)
        self.gmf_item_embed = nn.Embedding(n_items, embed_dim)

        # MLP embeddings
        self.mlp_user_embed = nn.Embedding(n_users, embed_dim)
        self.mlp_item_embed = nn.Embedding(n_items, embed_dim)

        # MLP layers
        mlp_modules: list[nn.Module] = []
        input_dim = embed_dim * 2
        for hidden_dim in mlp_layers:
            mlp_modules.append(nn.Linear(input_dim, hidden_dim))
            mlp_modules.append(nn.ReLU())
            mlp_modules.append(nn.Dropout(dropout))
            input_dim = hidden_dim
        self.mlp = nn.Sequential(*mlp_modules)

        # Final prediction layer (GMF output + last MLP layer output)
        self.output_layer = nn.Linear(embed_dim + mlp_layers[-1], 1)
        self.sigmoid = nn.Sigmoid()

        self._init_weights()

    def _init_weights(self) -> None:
        for module in self.modules():
            if isinstance(module, nn.Embedding):
                nn.init.normal_(module.weight, std=0.01)
            elif isinstance(module, nn.Linear):
                nn.init.xavier_uniform_(module.weight)
                if module.bias is not None:
                    nn.init.zeros_(module.bias)

    def forward(self, user_ids: torch.Tensor, item_ids: torch.Tensor) -> torch.Tensor:
        # GMF path
        gmf_user = self.gmf_user_embed(user_ids)
        gmf_item = self.gmf_item_embed(item_ids)
        gmf_out = gmf_user * gmf_item  # element-wise product

        # MLP path
        mlp_user = self.mlp_user_embed(user_ids)
        mlp_item = self.mlp_item_embed(item_ids)
        mlp_input = torch.cat([mlp_user, mlp_item], dim=-1)
        mlp_out = self.mlp(mlp_input)

        # Combine and predict
        combined = torch.cat([gmf_out, mlp_out], dim=-1)
        logit = self.output_layer(combined)
        return self.sigmoid(logit).squeeze(-1) * 9.0 + 1.0


# ---------------------------------------------------------------------------
# NCF training
# ---------------------------------------------------------------------------

MIN_RATINGS_FOR_NCF = 50


def train_ncf(
    ratings: list[dict[str, Any]],
    output_dir: Path,
    device: str,
    epochs: int,
    batch_size: int,
) -> None:
    """Train a Neural Collaborative Filtering model and save to disk.

    Saves:
        - ``ncf_model.pt``         -- model state_dict
        - ``ncf_metadata.json``    -- id mappings + hyper-parameters
    """
    print(f"\n{'='*60}")
    print("Training NCF model")
    print(f"{'='*60}")

    if len(ratings) < MIN_RATINGS_FOR_NCF:
        print(
            f"Skipping NCF: only {len(ratings)} ratings "
            f"(minimum {MIN_RATINGS_FOR_NCF} required)"
        )
        return

    t0 = time.time()

    # ── Build ID mappings ────────────────────────────────────────────────
    user_ids_set = sorted({r["user_id"] for r in ratings})
    movie_ids_set = sorted({r["movie_id"] for r in ratings})

    user_to_idx = {uid: idx for idx, uid in enumerate(user_ids_set)}
    movie_to_idx = {mid: idx for idx, mid in enumerate(movie_ids_set)}

    n_users = len(user_ids_set)
    n_items = len(movie_ids_set)
    print(f"Users: {n_users}  |  Items: {n_items}  |  Ratings: {len(ratings)}")

    # ── Prepare tensors ──────────────────────────────────────────────────
    user_indices = torch.tensor(
        [user_to_idx[r["user_id"]] for r in ratings], dtype=torch.long
    )
    item_indices = torch.tensor(
        [movie_to_idx[r["movie_id"]] for r in ratings], dtype=torch.long
    )
    rating_values = torch.tensor(
        [r["rating"] for r in ratings], dtype=torch.float32
    )

    # ── Train / validation split (90/10) ─────────────────────────────────
    n_total = len(ratings)
    perm = torch.randperm(n_total)
    split = int(n_total * 0.9)

    train_idx, val_idx = perm[:split], perm[split:]

    train_ds = TensorDataset(
        user_indices[train_idx], item_indices[train_idx], rating_values[train_idx]
    )
    val_ds = TensorDataset(
        user_indices[val_idx], item_indices[val_idx], rating_values[val_idx]
    )

    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True)
    val_loader = DataLoader(val_ds, batch_size=batch_size)

    print(f"Train: {len(train_ds)}  |  Val: {len(val_ds)}")

    # ── Instantiate model ────────────────────────────────────────────────
    model = NCFModel(n_users=n_users, n_items=n_items).to(device)
    criterion = nn.MSELoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

    total_params = sum(p.numel() for p in model.parameters())
    print(f"Model parameters: {total_params:,}")

    # ── Training loop with early stopping ────────────────────────────────
    best_val_loss = float("inf")
    patience = 5
    patience_counter = 0
    best_state = None

    for epoch in range(1, epochs + 1):
        # Train
        model.train()
        train_loss_sum = 0.0
        train_count = 0
        for u_batch, i_batch, r_batch in train_loader:
            u_batch = u_batch.to(device)
            i_batch = i_batch.to(device)
            r_batch = r_batch.to(device)

            optimizer.zero_grad()
            predictions = model(u_batch, i_batch)
            loss = criterion(predictions, r_batch)
            loss.backward()
            optimizer.step()

            train_loss_sum += loss.item() * len(r_batch)
            train_count += len(r_batch)

        train_loss = train_loss_sum / train_count

        # Validate
        model.eval()
        val_loss_sum = 0.0
        val_count = 0
        with torch.no_grad():
            for u_batch, i_batch, r_batch in val_loader:
                u_batch = u_batch.to(device)
                i_batch = i_batch.to(device)
                r_batch = r_batch.to(device)

                predictions = model(u_batch, i_batch)
                loss = criterion(predictions, r_batch)
                val_loss_sum += loss.item() * len(r_batch)
                val_count += len(r_batch)

        val_loss = val_loss_sum / max(val_count, 1)

        # Log every epoch (useful on GPU runs where you tail the log)
        print(
            f"  Epoch {epoch:3d}/{epochs}  "
            f"train_loss={train_loss:.4f}  val_loss={val_loss:.4f}"
        )

        # Early stopping check
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            patience_counter = 0
            best_state = {k: v.cpu().clone() for k, v in model.state_dict().items()}
        else:
            patience_counter += 1
            if patience_counter >= patience:
                print(f"  Early stopping at epoch {epoch} (patience={patience})")
                break

    # ── Save model ───────────────────────────────────────────────────────
    if best_state is None:
        print("Warning: no checkpoint saved (training may have failed)")
        return

    model_path = output_dir / "ncf_model.pt"
    torch.save(best_state, model_path)
    print(f"Saved best model (val_loss={best_val_loss:.4f}) to {model_path}")

    metadata = {
        "n_users": n_users,
        "n_items": n_items,
        "embed_dim": 32,
        "mlp_layers": [128, 64, 32],
        "dropout": 0.2,
        "best_val_loss": float(best_val_loss),
        "user_id_to_idx": {str(k): v for k, v in user_to_idx.items()},
        "movie_id_to_idx": {str(k): v for k, v in movie_to_idx.items()},
    }
    metadata_path = output_dir / "ncf_metadata.json"
    with open(metadata_path, "w") as f:
        json.dump(metadata, f, indent=2)
    print(f"Saved metadata to {metadata_path}")

    elapsed = time.time() - t0
    print(f"NCF training done in {elapsed:.1f}s")


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main() -> None:
    parser = argparse.ArgumentParser(
        description="Train movie-recommendation models on GPU (RunPod)."
    )
    parser.add_argument(
        "--data-dir",
        type=Path,
        default=Path(__file__).resolve().parent.parent / "data" / "export",
        help="Directory containing movies.json and ratings.json",
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=Path(__file__).resolve().parent.parent / "data" / "models",
        help="Directory to write trained model artefacts",
    )
    parser.add_argument(
        "--device",
        type=str,
        default="cuda:0",
        help="Torch device (default: cuda:0)",
    )
    parser.add_argument(
        "--epochs",
        type=int,
        default=50,
        help="Max training epochs for NCF (default: 50)",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=256,
        help="Batch size for NCF training (default: 256)",
    )
    args = parser.parse_args()

    # ── Resolve paths ────────────────────────────────────────────────────
    data_dir: Path = args.data_dir
    output_dir: Path = args.output_dir
    output_dir.mkdir(parents=True, exist_ok=True)

    # ── Validate device ──────────────────────────────────────────────────
    device = args.device
    if device.startswith("cuda") and not torch.cuda.is_available():
        print(f"Warning: {device} requested but CUDA is not available, falling back to cpu")
        device = "cpu"
    print(f"Using device: {device}")

    # ── Load data ────────────────────────────────────────────────────────
    movies_path = data_dir / "movies.json"
    ratings_path = data_dir / "ratings.json"

    if not movies_path.exists():
        print(f"Error: {movies_path} not found. Run export_data.py first.")
        raise SystemExit(1)

    with open(movies_path) as f:
        movies: list[dict[str, Any]] = json.load(f)

    ratings: list[dict[str, Any]] = []
    if ratings_path.exists():
        with open(ratings_path) as f:
            ratings = json.load(f)
    else:
        print(f"Warning: {ratings_path} not found, skipping NCF training")

    print(f"\nLoaded {len(movies)} movies, {len(ratings)} ratings from {data_dir}")

    # ── Train ────────────────────────────────────────────────────────────
    overall_start = time.time()

    train_embeddings(movies, output_dir, device)
    train_ncf(ratings, output_dir, device, args.epochs, args.batch_size)

    total = time.time() - overall_start
    print(f"\nAll done in {total:.1f}s")


if __name__ == "__main__":
    main()