"""Build the multi-criteria rating table used by Nassar et al. (2020) from the LARA TripAdvisor files.

Source: Hongning Wang's LARA TripAdvisor "Review_Texts" (one hotel_<id>.dat file per hotel), archived at
http://web.archive.org/web/20150316061732/http://times.cs.uiuc.edu/~wang296/Data/LARA/TripAdvisor/Review_Texts.zip

Filtering (Section 4.1 of the paper only says the data has 7 criteria and an overall rating on 1-5):
  1. keep reviews whose overall rating and all 7 criteria ratings are present and in 1..5
     (the files use -1 or 0 for a missing rating);
  2. drop reviews with no author name, or the anonymous "A TripAdvisor Member";
  3. keep one rating per (user, hotel) pair: the first one in the file.

Usage:  python prepare_data.py [raw_dir] [out_csv]
"""
from __future__ import annotations

import re
import sys
import zipfile
from pathlib import Path

import pandas as pd

CRITERIA = ["value", "rooms", "location", "cleanliness", "check_in", "service", "business_service"]
FIELDS = {  # tag in the .dat file -> column name
    "Overall": "overall",
    "Value": "value",
    "Rooms": "rooms",
    "Location": "location",
    "Cleanliness": "cleanliness",
    "Check in / front desk": "check_in",
    "Service": "service",
    "Business service": "business_service",
}
TAG = re.compile(r"^<(" + "|".join(re.escape(k) for k in FIELDS) + r")>\s*(-?\d+(?:\.\d+)?)\s*$")
ANONYMOUS = {"", "a tripadvisor member"}


def parse_hotel_file(name: str, text: str) -> list[dict]:
    """Return one dict per review: user, hotel and the 8 ratings (None where missing)."""
    hotel = re.search(r"hotel_(\d+)", name).group(1)
    reviews, cur = [], None
    for line in text.splitlines():
        if line.startswith("<Author>"):
            if cur is not None:
                reviews.append(cur)
            cur = {"user": line[len("<Author>"):].strip(), "hotel": hotel}
            continue
        if cur is None:
            continue
        m = TAG.match(line.strip())
        if m:
            cur[FIELDS[m.group(1)]] = float(m.group(2))
    if cur is not None:
        reviews.append(cur)
    return reviews


def read_hotel_files(raw_dir: Path) -> list[tuple[str, str]]:
    """(file name, text) for every hotel_*.dat, read from Review_Texts.zip or from unzipped files."""
    decode = lambda b: b.decode("utf-8", errors="replace")
    zips = sorted(raw_dir.rglob("Review_Texts.zip"))
    if zips:
        with zipfile.ZipFile(zips[0]) as z:
            return [(Path(n).name, decode(z.read(n))) for n in sorted(z.namelist())
                    if re.search(r"hotel_\d+\.dat$", n)]
    return [(p.name, decode(p.read_bytes())) for p in sorted(raw_dir.rglob("hotel_*.dat"))]


def build_table(raw_dir: Path) -> tuple[pd.DataFrame, dict]:
    files = read_hotel_files(raw_dir)
    if not files:
        raise FileNotFoundError(f"no Review_Texts.zip or hotel_*.dat files under {raw_dir}")
    df = pd.DataFrame([r for name, text in files for r in parse_hotel_file(name, text)])
    cols = ["overall"] + CRITERIA
    stats = {"hotel_files": len(files), "reviews_parsed": len(df)}

    ok = df[cols].notna().all(axis=1) & df[cols].isin([1.0, 2.0, 3.0, 4.0, 5.0]).all(axis=1)
    df = df[ok]
    stats["complete_8_ratings"] = len(df)

    df = df[~df["user"].str.lower().isin(ANONYMOUS)]
    stats["after_dropping_anonymous"] = len(df)

    df = df.drop_duplicates(["user", "hotel"], keep="first")
    stats["after_dedup_user_hotel"] = len(df)

    df[cols] = df[cols].astype(int)
    return df[["user", "hotel"] + cols].reset_index(drop=True), stats


def describe(df: pd.DataFrame) -> tuple[pd.DataFrame, pd.DataFrame]:
    """Tables 1 and 2 of the paper, for our copy of the data, next to the paper's numbers."""
    n_u, n_i, n = df.user.nunique(), df.hotel.nunique(), len(df)
    t1 = pd.DataFrame(
        {
            "paper": [72119, 1850, 81085, "99.94%"],
            "ours": [n_u, n_i, n, f"{100 * (1 - n / (n_u * n_i)):.2f}%"],
        },
        index=["users", "hotels", "ratings", "sparsity"],
    )
    paper_t2 = {
        "value": [4268, 5853, 10374, 24696, 35894],
        "rooms": [4002, 5517, 10493, 25019, 36054],
        "location": [1464, 3141, 8503, 22177, 45800],
        "cleanliness": [2591, 3585, 8232, 19951, 46726],
        "check_in": [3775, 4468, 11083, 19201, 42558],
        "service": [4086, 4462, 10036, 20114, 42387],
        "business_service": [4525, 5849, 19624, 21661, 29426],
        "overall": [4246, 6052, 8078, 24180, 38529],
    }
    rows = []
    for c, paper in paper_t2.items():
        ours = df[c].value_counts().reindex([1, 2, 3, 4, 5], fill_value=0).tolist()
        rows.append([c, "paper", *paper])
        rows.append([c, "ours", *ours])
    t2 = pd.DataFrame(rows, columns=["rating", "source", "1", "2", "3", "4", "5"])
    return t1, t2


if __name__ == "__main__":
    here = Path(__file__).resolve().parent
    raw = Path(sys.argv[1]) if len(sys.argv) > 1 else here / "data" / "raw"
    out = Path(sys.argv[2]) if len(sys.argv) > 2 else here / "data" / "tripadvisor_mc.csv"
    table, stats = build_table(raw)
    table.to_csv(out, index=False)
    for k, v in stats.items():
        print(f"{k:28s} {v:>8,}")
    t1, t2 = describe(table)
    print("\nTable 1\n", t1.to_string(), "\n\nTable 2\n", t2.to_string(index=False), sep="")
    print(f"\nwrote {out}")