src/loader.py

"""
Loader for the BookRec recommender on the Goodbooks-10k dataset.

Loads the COMPLETE real dataset:
  books.csv     -> 10,000 books (title, authors, isbn, live image_url)
  ratings.csv   -> 5,976,479 explicit 1-5 ratings (user_id, book_id, rating)
  book_tags.csv -> book -> tag_id with counts (the genre/content signal)
  tags.csv      -> tag_id -> tag_name

Design decisions handled here:
  * Item key is ISBN = str(book_id), so the pipeline keys on one string.
  * Genres are the top tags per book (shelf-labelled stop-tags removed).
  * Covers prefer the live Goodreads (gr-assets) https URL, with an Open
    Library fallback by ISBN, so a cover URL is always available.
"""
from __future__ import annotations

from pathlib import Path
from typing import Optional, Tuple

import numpy as np
import pandas as pd

DATA_DIR = Path(__file__).resolve().parents[1] / "data"


# ---------------------------------------------------------------------------
# Covers (resilient)
# ---------------------------------------------------------------------------

def cover_url(isbn: str, size: str = "M") -> str:
    """Best-effort cover URL: Open Library API (free, no key, live)."""
    isbn = str(isbn).strip()
    if not isbn:
        return ""
    size_map = {"s": "S", "m": "M", "l": "L"}
    letter = size_map.get(str(size).lower(), "M")
    return f"https://covers.openlibrary.org/b/isbn/{isbn}-{letter}.jpg"


def pick_cover_url(row: pd.Series) -> str:
    """Return the best available cover URL for a book row.

    Prefers a live HTTPS cover (Goodreads gr-assets), else Open Library by
    ISBN.
    """
    for col in ("Image-URL-M", "Image-URL-L", "Image-URL-S"):
        v = row.get(col)
        if isinstance(v, str) and v.startswith("https://"):
            return v
    isbn = str(row.get("ISBN", "")).strip()
    if isbn:
        return cover_url(isbn, "M")
    return ""


def isbn_cover(isbn: str) -> str:
    return cover_url(isbn, "M")


# ---------------------------------------------------------------------------
# Goodbooks-10k loader (ratings + genres/tags + live covers)
# ---------------------------------------------------------------------------

GB_STOP_TAGS = {"to-read", "currently-reading", "favorites", "books-i-own",
                "read", "owned", "default", "re-read", "favorites2"}


def load_goodbooks() -> Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame,
                              pd.DataFrame]:
    """Load the four Goodbooks-10k raw files."""
    d = DATA_DIR
    books = pd.read_csv(d / "books.csv")
    ratings = pd.read_csv(d / "ratings.csv")
    book_tags = pd.read_csv(d / "book_tags.csv")
    tags = pd.read_csv(d / "tags.csv")
    return books, ratings, book_tags, tags


def build_genres(books, book_tags, tags, top_k: int = 15) -> pd.Series:
    """Map each goodreads_book_id to a space-joined string of its top genres.

    Joins book_tags -> tags, drops shelf-labelled stop-tags, keeps the top-K
    highest-count genre tags per book as a content string.
    """
    m = book_tags.merge(tags, on="tag_id")
    m = m[~m["tag_name"].astype(str).isin(GB_STOP_TAGS)]

    def top(db):
        g = db.sort_values("count", ascending=False).head(top_k)
        return " ".join(g["tag_name"].astype(str))

    genres = m.groupby("goodreads_book_id").apply(top, include_groups=False)
    genres = genres.rename("genres_text")
    return genres


def prepare_goodbooks() -> dict:
    """Load Goodbooks-10k and normalise to the pipeline's internal schema.

    Returns books/ratings with the standard columns the recommender expects
    (ISBN = str(book_id), Book-Title, Book-Author, Image-URL-M, Book-Genres)
    plus a genres-text channel for content similarity.
    """
    books, rat, book_tags, tags = load_goodbooks()

    genres = build_genres(books, book_tags, tags, top_k=15)

    books = books.copy()
    books["ISBN"] = books["book_id"].astype(str)
    books["Book-Title"] = books["title"]
    books["Book-Author"] = books["authors"]
    books["Image-URL-M"] = books["image_url"]
    books["Book-Genres"] = (books["goodreads_book_id"].map(genres)
                             .fillna(""))

    rat = rat.rename(columns={"user_id": "User-ID", "book_id": "ISBN",
                              "rating": "Book-Rating"})
    rat["User-ID"] = rat["User-ID"].astype(str)
    rat["ISBN"] = rat["ISBN"].astype(str)

    return {
        "books": books.reset_index(drop=True),
        "ratings": rat.reset_index(drop=True),
        "n_books": len(books),
        "n_ratings": len(rat),
        "n_users": rat["User-ID"].nunique(),
        "explicit_share": 1.0,  # Goodbooks ratings are all explicit (1-5)
        "implicit_share": 0.0,
        "source": "goodbooks-10k",
    }


def prepare_amazon() -> dict:
    """Load the sampled Amazon Books parquets into the pipeline schema.

    Reads amazon_books/books.parquet and amazon_books/ratings.parquet
    (converted from the McAuley-Lab/Amazon-Reviews-2023 benchmark) and maps
    them to the standard ISBN / Book-Title / Book-Author / Book-Genres /
    User-ID / Book-Rating columns the recommender expects.
    """
    import pandas as pd

    base = Path(__file__).resolve().parents[1] / "amazon_books"
    books = pd.read_parquet(base / "books.parquet")
    rat = pd.read_parquet(base / "ratings.parquet")

    def _flatten(v):
        if v is None or isinstance(v, str) or not hasattr(v, "__iter__"):
            return [v]
        out = []
        for x in v:
            out.extend(_flatten(x))
        return out

    def _first_author(a) -> str:
        if a is None or isinstance(a, str):
            return "" if a is None else a
        flat = _flatten(a)
        return str(flat[0]) if flat else ""

    def _genres(c) -> str:
        flat = _flatten(c)
        cleaned = [str(x) for x in flat if str(x).lower() != "books"]
        return " ".join(dict.fromkeys(cleaned))

    books = books.copy()
    books["ISBN"] = books["item_id"].astype(str)
    books["Book-Title"] = books["title"].fillna("")
    books["Book-Author"] = (books["authors"].map(_first_author)
                            .fillna(""))
    books["Book-Genres"] = (books["categories"].map(_genres)
                            .fillna(""))

    rat = rat.rename(columns={"user_id": "User-ID", "item_id": "ISBN",
                              "rating": "Book-Rating"})
    rat["User-ID"] = rat["User-ID"].astype(str)
    rat["ISBN"] = rat["ISBN"].astype(str)

    return {
        "books": books.reset_index(drop=True),
        "ratings": rat.reset_index(drop=True),
        "n_books": len(books),
        "n_ratings": len(rat),
        "n_users": rat["User-ID"].nunique(),
        "explicit_share": 1.0,
        "implicit_share": 0.0,
        "source": "amazon-books-5core",
    }