import os
import random

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

import torch
import torch.nn as nn
import torch.nn.functional as F

from typing import Tuple, Callable

from sklearn.model_selection import train_test_split
from sklearn.metrics import root_mean_squared_error

from torch_geometric.nn import LGConv
from torch_sparse import SparseTensor

# Constants
DATA_DIR = ""
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

SEED = 42
torch.manual_seed(SEED)
np.random.seed(SEED)

def read_data_df() -> Tuple[pd.DataFrame, pd.DataFrame]:
    """Reads in data and splits it into training and validation sets with a 75/25 split."""

    df = pd.read_csv(os.path.join(DATA_DIR, "train_ratings.csv"))

    # Split sid_pid into sid and pid columns
    df[["sid", "pid"]] = df["sid_pid"].str.split("_", expand=True)
    df = df.drop("sid_pid", axis=1)
    df["sid"] = df["sid"].astype(int)
    df["pid"] = df["pid"].astype(int)

    # Split into train and validation dataset
    train_df, valid_df = train_test_split(df, test_size=0.1)
    return train_df, valid_df


def read_data_matrix(df: pd.DataFrame) -> np.ndarray:
    """Returns matrix view of the training data, where columns are scientists (sid) and
    rows are papers (pid)."""

    return df.pivot(index="sid", columns="pid", values="rating").values


def read_wishlist_df() -> pd.DataFrame:
    """Reads in the wishlist data (train_tbr.csv)."""

    tbr_df = pd.read_csv(os.path.join(DATA_DIR, "train_tbr.csv"))
    # Ensure sid and pid are integers
    tbr_df["sid"] = tbr_df["sid"].astype(int)
    tbr_df["pid"] = tbr_df["pid"].astype(int)
    return tbr_df


def read_wishlist_matrix(df: pd.DataFrame) -> np.ndarray:
    """Returns matrix view of the wishlist data.
    Rows are scientists (sid), columns are papers (pid).
    Values are 1 if the paper is on the wishlist, NaN otherwise."""

    # Add a temporary column with value 1 for pivoting
    df_copy = df.copy()
    df_copy['wishlisted'] = 1


    all_sids = pd.Index(range(10000))
    wishlist_matrix = df_copy.pivot(index="sid", columns="pid", values="wishlisted")
    wishlist_matrix = wishlist_matrix.reindex(all_sids)

    return wishlist_matrix.values

def evaluate(valid_df: pd.DataFrame, pred_fn: Callable[[np.ndarray, np.ndarray], np.ndarray]) -> float:
    """
    Inputs:
        valid_df: Validation data, returned from read_data_df for example.
        pred_fn: Function that takes in arrays of sid and pid and outputs their rating predictions.

    Outputs: Validation RMSE
    """

    preds = pred_fn(valid_df["sid"].values, valid_df["pid"].values)
    return root_mean_squared_error(valid_df["rating"].values, preds)


def make_submission(pred_fn: Callable[[np.ndarray, np.ndarray], np.ndarray], filename: os.PathLike):
    """Makes a submission CSV file that can be submitted to kaggle.

    Inputs:
        pred_fn: Function that takes in arrays of sid and pid and outputs a score.
        filename: File to save the submission to.
    """

    df = pd.read_csv(os.path.join(DATA_DIR, "sample_submission.csv"))

    # Get sids and pids
    sid_pid = df["sid_pid"].str.split("_", expand=True)
    sids = sid_pid[0]
    pids = sid_pid[1]
    sids = sids.astype(int).values
    pids = pids.astype(int).values

    df["rating"] = pred_fn(sids, pids)
    df.to_csv(filename, index=False)

def impute_values(mat: np.ndarray) -> np.ndarray:
    """Fills in missing values of the matrix"""
    return np.nan_to_num(mat, nan=3.0)

train_ratings_df, valid_ratings_df = read_data_df()
train_wishlist_df = read_wishlist_df()

train_mat = read_data_matrix(train_ratings_df)
train_mat = impute_values(train_mat)

train_wishlist_mat = read_wishlist_matrix(train_wishlist_df)

singular_values = np.linalg.svd(train_mat, compute_uv=False, hermitian=False)
plt.plot(singular_values)
plt.xlabel("Index")
plt.ylabel("Value")
plt.title("Singular value spectrum")
plt.show()

def opt_rank_k_approximation(m: np.ndarray, k: int):
    """Returns the optimal rank-k reconstruction matrix, using SVD."""

    assert 0 < k <= np.min(m.shape), f"The rank must be in [0, min(m, n)]"

    U, S, Vh = np.linalg.svd(m, full_matrices=False)

    U_k = U[:, :k]
    S_k = S[:k]
    Vh_k = Vh[:k]

    return np.dot(U_k * S_k, Vh_k)


def matrix_pred_fn(train_recon: np.ndarray, sids: np.ndarray, pids: np.ndarray) -> np.ndarray:
    """
    Input:
        train_recon: (M, N) matrix with predicted values for every (sid, pid) pair.
        sids: (D,) vector with integer scientist IDs.
        pids: (D,) vector with integer paper IDs.

    Outputs: (D,) vector with predictions.
    """

    return train_recon[sids, pids]

train_recon = opt_rank_k_approximation(train_mat, k=2)

pred_fn = lambda sids, pids: matrix_pred_fn(train_recon, sids, pids)

# Evaluate on validation data
val_score = evaluate(valid_ratings_df, pred_fn)
print(f"Validation RMSE: {val_score:.3f}")

class EmbeddingDotProductModel(nn.Module):
    def __init__(self, num_scientists: int, num_papers: int, dim: int):
        super().__init__()

        # Assign to each scientist and paper an embedding
        self.scientist_emb = nn.Embedding(num_scientists, dim)
        self.paper_emb = nn.Embedding(num_papers, dim)

    def forward(self, sid: torch.Tensor, pid: torch.Tensor) -> torch.Tensor:
        """
        Inputs:
            sid: [B,], int
            pid: [B,], int

        Outputs: [B,], float
        """

        # Per-pair dot product
        return torch.sum(self.scientist_emb(sid) * self.paper_emb(pid), dim=-1)

# Define model (10k scientists, 1k papers, 32-dimensional embeddings) and optimizer
model = EmbeddingDotProductModel(10_000, 1_000, 64).to(device)
optim = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-5)

def get_dataset(df: pd.DataFrame):
    sids = torch.tensor(df["sid"].values, dtype=torch.long)
    pids = torch.tensor(df["pid"].values, dtype=torch.long)
    ratings = torch.tensor(df["rating"].values, dtype=torch.float32)
    return torch.utils.data.TensorDataset(sids, pids, ratings)

train_dataset = get_dataset(train_ratings_df)
valid_dataset = get_dataset(valid_ratings_df)

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)
valid_loader = torch.utils.data.DataLoader(valid_dataset, batch_size=64, shuffle=False)

NUM_EPOCHS = 5
for epoch in range(NUM_EPOCHS):
    model.train()
    total_loss, total_count = 0.0, 0
    for sid, pid, ratings in train_loader:
        sid = sid.to(device)
        pid = pid.to(device)
        ratings = ratings.to(device)

        preds = model(sid, pid)
        loss = F.mse_loss(preds, ratings)

        optim.zero_grad()
        loss.backward()
        optim.step()

        total_loss += loss.item() * len(sid)
        total_count += len(sid)

    print(f"[Epoch {epoch+1}/{NUM_EPOCHS}] Train Loss: {total_loss / total_count:.4f}")

pred_fn = lambda sids, pids: model(torch.from_numpy(sids).to(device), torch.from_numpy(pids).to(device))

# Evaluate on validation data
with torch.no_grad():
    val_score = evaluate(valid_ratings_df, pred_fn)

print(f"Validation RMSE: {val_score:.3f}")

def iterative_svd(df, k=2, max_iterations=10, tolerance=0.001):
    mat = read_data_matrix(df)
    missing_mask = np.isnan(mat)
    imputed_mat = impute_values(mat)

    for iteration in range(max_iterations):
        previous_mat = imputed_mat.copy()
        reconstructed_mat = opt_rank_k_approximation(imputed_mat, k)

        imputed_mat = mat.copy()
        imputed_mat[missing_mask] = reconstructed_mat[missing_mask]

        diff = np.linalg.norm(reconstructed_mat[missing_mask] - previous_mat[missing_mask])
        print(f"Iteration {iteration + 1}, Change: {diff:.6f}")

        if diff < tolerance:
            print(f"Converged after {iteration + 1} iterations")
            break

    return imputed_mat

final_mat = iterative_svd(train_ratings_df, k=2, max_iterations=200, tolerance=0.0001)

pred_iter_svd = lambda sids, pids: matrix_pred_fn(final_mat, sids, pids).flatten()

val_score = evaluate(valid_ratings_df, pred_iter_svd)
print(f"Validation RMSE: {val_score:.3f}")

class SVDPP(nn.Module):
    def __init__(self, num_users, num_items, emb_dim=32):
        super().__init__()
        self.user_emb = nn.Embedding(num_users, emb_dim)
        self.item_emb = nn.Embedding(num_items, emb_dim)
        self.user_bias = nn.Embedding(num_users, 1)
        self.item_bias = nn.Embedding(num_items, 1)
        self.implicit_emb = nn.Embedding(num_items + 1, emb_dim, padding_idx=num_items)  # Padding index = num_items
        self.global_mean = nn.Parameter(torch.tensor(3.0), requires_grad=False)

    def forward(self, user_ids, item_ids, implicit_items):
        p_u = self.user_emb(user_ids)
        q_i = self.item_emb(item_ids)
        b_u = self.user_bias(user_ids).squeeze()
        b_i = self.item_bias(item_ids).squeeze()

        mask = (implicit_items != self.implicit_emb.padding_idx).unsqueeze(-1)
        y_j = self.implicit_emb(implicit_items) * mask
        norm = torch.sqrt(mask.sum(dim=1).clamp(min=1).float())
        implicit_sum = y_j.sum(dim=1) / norm

        x_u = p_u + implicit_sum
        dot = (x_u * q_i).sum(dim=1)

        return self.global_mean + b_u + b_i + dot

model_svdpp = SVDPP(num_users=10_000, num_items=1_000, emb_dim=16).to(device)
optim_svdpp = torch.optim.Adam(model_svdpp.parameters(), lr=1e-3, weight_decay=1e-4)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optim_svdpp, mode='min', factor=0.5, patience=1)

from collections import defaultdict

def build_implicit_item_matrix(df, num_users, num_items, max_len=50):
    PAD_IDX = num_items
    user_implicit = defaultdict(list)
    for _, row in df.iterrows():
        user_implicit[int(row['sid'])].append(int(row['pid']))
    matrix = np.full((num_users, max_len), fill_value=PAD_IDX, dtype=int)
    for user_id, items in user_implicit.items():
        matrix[user_id, :min(len(items), max_len)] = items[:max_len]
    return torch.tensor(matrix, dtype=torch.long)

NUM_USERS = 10_000
NUM_ITEMS = 1_000
EMB_DIM = 16
MAX_LEN = 50

implicit_matrix = build_implicit_item_matrix(train_ratings_df, NUM_USERS, NUM_ITEMS, max_len=MAX_LEN).to(device)

# Cell 4: Training and Validation Loop with Masking Current Item and RMSE Logging
NUM_EPOCHS = 10
PAD_IDX = model_svdpp.implicit_emb.padding_idx

for epoch in range(NUM_EPOCHS):
    model_svdpp.train()
    train_loss_total, train_count = 0.0, 0
    for sid, pid, ratings in train_loader:
        sid = sid.to(device)
        pid = pid.to(device)
        ratings = ratings.to(device)

        # clone and mask out current item from implicit feedback
        imp_items = implicit_matrix[sid].clone()
        imp_items[imp_items == pid.unsqueeze(1)] = PAD_IDX

        preds = model_svdpp(sid, pid, imp_items)
        loss = F.mse_loss(preds, ratings)

        optim_svdpp.zero_grad()
        loss.backward()
        optim_svdpp.step()

        train_loss_total += loss.item() * sid.size(0)
        train_count += sid.size(0)

    avg_train_loss = train_loss_total / train_count

    model_svdpp.eval()
    val_mse_total, val_count = 0.0, 0
    with torch.no_grad():
        for sid, pid, ratings in valid_loader:
            sid = sid.to(device)
            pid = pid.to(device)
            ratings = ratings.to(device)

            imp_items = implicit_matrix[sid].clone()
            imp_items[imp_items == pid.unsqueeze(1)] = PAD_IDX

            preds = model_svdpp(sid, pid, imp_items)
            mse = F.mse_loss(preds, ratings, reduction='sum')

            val_mse_total += mse.item()
            val_count += sid.size(0)

    valid_rmse = (val_mse_total / val_count) ** 0.5
    print(f"Epoch {epoch + 1}/{NUM_EPOCHS} – Train Loss: {avg_train_loss:.4f}, Valid RMSE: {valid_rmse:.4f}")
    scheduler.step(val_mse_total / val_count)

pred_fn_svdpp = lambda sids, pids: model_svdpp(
    torch.from_numpy(sids).to(device),
    torch.from_numpy(pids).to(device),
    implicit_matrix[
        torch.from_numpy(np.clip(sids, 0, implicit_matrix.shape[0] - 1)).to(device)
    ]
).detach().cpu().numpy()

with torch.no_grad():
    val_score = evaluate(valid_ratings_df, pred_fn_svdpp)

print(f"Validation RMSE: {val_score:.3f}")

def get_dataset(
    df: pd.DataFrame,
    wishlist_df: pd.DataFrame = None
) -> torch.utils.data.Dataset:
    """
    Returns a dataset of (sid, pid, rating, wishlist_flag) for every
    (sid,pid) in ratings OR wishlist.  Missing ratings or wishlist
    entries are filled with 0.0.
    """
    # If no wishlist provided, just add an empty wishlist column
    if wishlist_df is None:
        merged = df.copy()
        merged['wishlist'] = 0.0

    else:
        w = wishlist_df[['sid', 'pid']].copy()
        w['wishlist'] = 1.0

        merged = pd.merge(
            df[['sid', 'pid', 'rating']],
            w,
            on=['sid', 'pid'],
            how='outer'       # <-- full join
        )

        merged['rating'] = merged['rating'].fillna(0.0)
        merged['wishlist'] = merged['wishlist'].fillna(0.0)

    sids      = torch.from_numpy(merged['sid'].to_numpy()).long()
    pids      = torch.from_numpy(merged['pid'].to_numpy()).long()
    ratings   = torch.from_numpy(merged['rating'].to_numpy()).float()
    wishlists = torch.from_numpy(merged['wishlist'].to_numpy()).float()

    return torch.utils.data.TensorDataset(sids, pids, ratings, wishlists)

class GMF(nn.Module):
    def __init__(self, num_researchers, num_papers, emb_dim):
        super().__init__()
        self.researcher_emb = nn.Embedding(num_researchers, emb_dim)
        self.paper_emb = nn.Embedding(num_papers, emb_dim)
        self.dropout = nn.Dropout(0.2)

    def forward(self, researcher_indices, paper_indices):
        r = self.dropout(self.researcher_emb(researcher_indices))
        p = self.dropout(self.paper_emb(paper_indices))
        return r * p

class MLP(nn.Module):
    def __init__(self, num_researchers, num_papers, emb_dim, mlp_layers):
        super().__init__()
        self.researcher_emb = nn.Embedding(num_researchers, emb_dim)
        self.paper_emb = nn.Embedding(num_papers, emb_dim)
        self.dropout = nn.Dropout(0.2)
        layers = []
        input_size = emb_dim * 2
        for h in mlp_layers:
            layers += [nn.Linear(input_size, h), nn.GELU(), nn.Dropout(0.2)]
            input_size = h
        self.mlp = nn.Sequential(*layers)

    def forward(self, researcher_indices, paper_indices):
        r = self.dropout(self.researcher_emb(researcher_indices))
        p = self.dropout(self.paper_emb(paper_indices))
        x = torch.cat([r, p], dim=1)
        return self.mlp(x)

class NeuMF(nn.Module):
    def __init__(self, num_researchers, num_papers, emb_dim, mlp_layers, wishlist_weight=2.0):
        super().__init__()
        self.gmf = GMF(num_researchers, num_papers, emb_dim)
        self.mlp = MLP(num_researchers, num_papers, emb_dim, mlp_layers)
        self.alpha = nn.Parameter(torch.tensor(0.5))
        self.wishlist_weight = wishlist_weight

        final_size = emb_dim + mlp_layers[-1]
        self.fc_rating = nn.Linear(final_size, 1)
        self.fc_wishlist = nn.Linear(final_size, 1)

    def forward(self, researcher_indices, paper_indices):
        g = self.gmf(researcher_indices, paper_indices)
        m = self.mlp(researcher_indices, paper_indices)
        fused = torch.cat([self.alpha * g, (1 - self.alpha) * m], dim=1)
        rating = 1 + 4 * torch.sigmoid(self.fc_rating(fused).squeeze(1))
        wishlist = torch.sigmoid(self.fc_wishlist(fused).squeeze(1))
        return rating, wishlist

    def calc_loss(self, rating_pred, rating_target, wishlist_pred, wishlist_target):
        # Only compute rating loss where wishlist_target is 0
        mask = wishlist_target == 0
        rating_loss = F.mse_loss(rating_pred[mask], rating_target[mask])
        wishlist_loss = F.mse_loss(wishlist_pred, wishlist_target)
        return rating_loss + self.wishlist_weight * wishlist_loss

train_dataset = get_dataset(train_ratings_df, train_wishlist_df)
valid_dataset = get_dataset(valid_ratings_df)

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)
valid_loader = torch.utils.data.DataLoader(valid_dataset, batch_size=64, shuffle=False)

NUM_USERS = 10_000
NUM_ITEMS = 1_000
EMB_DIM = 16
MLP_HIDDEN = [32, 32]
WISHLIST_LOSS_WEIGHT = 2.0
BATCH_SIZE = 64
NUM_EPOCHS = 40
LEARNING_RATE = 5e-3

model_neumf = NeuMF(NUM_USERS, NUM_ITEMS, EMB_DIM, MLP_HIDDEN, WISHLIST_LOSS_WEIGHT).to(device)
optimizer = torch.optim.AdamW(model_neumf.parameters(), lr=LEARNING_RATE)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=1)

NUM_EPOCHS = 40
for epoch in range(NUM_EPOCHS):
    model_neumf.train()
    running_loss = 0.0

    for sid, pid, ratings, wishlists in train_loader:
        sid, pid = sid.to(device), pid.to(device)
        ratings, wishlists = ratings.to(device), wishlists.to(device)

        pred_ratings, pred_wishlist = model_neumf(sid, pid)
        loss = model_neumf.calc_loss(pred_ratings, ratings, pred_wishlist, wishlists)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        running_loss += loss.item() * sid.size(0)

    avg_train_loss = running_loss / len(train_loader.dataset)

    # Validation
    model_neumf.eval()
    mse_total = 0.0
    with torch.no_grad():
        for sid, pid, ratings, _ in valid_loader:
            sid, pid = sid.to(device), pid.to(device)
            ratings = ratings.to(device)

            pred_ratings, _ = model_neumf(sid, pid)
            mse = F.mse_loss(pred_ratings.clamp(1, 5), ratings)
            mse_total += mse.item() * sid.size(0)

    rmse = (mse_total / len(valid_loader.dataset)) ** 0.5
    print(f"Epoch {epoch + 1}/{NUM_EPOCHS} – Train Loss: {avg_train_loss:.4f}, Valid RMSE: {rmse:.4f}")
    scheduler.step(rmse)

os.makedirs("models", exist_ok=True)
torch.save(model_neumf.state_dict(), "models/neumf.pt")
print("Model saved to models/neumf.pt")

pred_fn_neumf = lambda sids, pids: model_neumf(
    torch.from_numpy(sids).to(device),
    torch.from_numpy(pids).to(device)
)[0].clamp(1, 5).detach().cpu().numpy()

# Evaluate on validation data
with torch.no_grad():
    val_score = evaluate(valid_ratings_df, pred_fn_neumf)

print(f"Validation RMSE: {val_score:.3f}")

class LightGCN(nn.Module):
    def __init__(self, num_users, num_items, emb_dim, num_layers=3):
        super().__init__()
        self.num_users  = num_users
        self.num_items  = num_items
        self.emb_dim    = emb_dim
        self.num_layers = num_layers

        self.user_emb = nn.Embedding(num_users, emb_dim)
        self.item_emb = nn.Embedding(num_items, emb_dim)
        nn.init.xavier_uniform_(self.user_emb.weight)
        nn.init.xavier_uniform_(self.item_emb.weight)

        # graph convolution layers
        from torch_geometric.nn import LGConv
        self.convs = nn.ModuleList([
            LGConv(normalize=True) for _ in range(num_layers)
        ])

        # will hold a SparseTensor
        self.adj = None

    def set_adj(self, edge_index: torch.LongTensor, device: torch.device):
        # build and store the adjacency once
        self.adj = SparseTensor(
            row=edge_index[0],
            col=edge_index[1] + self.num_users,
            sparse_sizes=(self.num_users + self.num_items,) * 2
        ).to(device)

    def forward(self):
        x = torch.cat([self.user_emb.weight, self.item_emb.weight], dim=0)
        embeddings = [x]

        for conv in self.convs:
            x = conv(x, self.adj)
            embeddings.append(x)

        x = torch.mean(torch.stack(embeddings), dim=0)
        user_emb, item_emb = torch.split(x, [self.num_users, self.num_items])
        return user_emb, item_emb

    def get_embeddings(self, user_idx, item_idx):
        return self.user_emb(user_idx), self.item_emb(item_idx)


class GraphNeuMF(nn.Module):
    def __init__(self, num_users, num_items, emb_dim, mlp_layers, wishlist_weight):
        super().__init__()
        self.lightgcn       = LightGCN(num_users, num_items, emb_dim)
        self.gmf            = GMF(num_users, num_items, emb_dim)
        self.mlp            = MLP(num_users, num_items, emb_dim, mlp_layers)
        self.alpha          = nn.Parameter(torch.tensor(0.5))
        self.wishlist_weight = wishlist_weight

        final_size = emb_dim + mlp_layers[-1] + emb_dim + emb_dim
        self.fc_rating  = nn.Linear(final_size, 1)
        self.fc_wishlist = nn.Linear(final_size, 1)

    def forward(self, user_idx, item_idx):
        user_emb_g, item_emb_g = self.lightgcn.get_embeddings(user_idx, item_idx)
        gmf_out  = self.gmf(user_idx, item_idx)
        mlp_out  = self.mlp(user_idx, item_idx)

        fused = torch.cat([
            self.alpha * gmf_out,
            (1.0 - self.alpha) * mlp_out,
            user_emb_g,
            item_emb_g
        ], dim=1)
        fused = fused if fused.dim() > 1 else fused.unsqueeze(0)

        rating_pred   = 1.0 + 4.0 * torch.sigmoid(self.fc_rating(fused).squeeze(1))
        wishlist_prob = torch.sigmoid(self.fc_wishlist(fused).squeeze(1))
        return rating_pred, wishlist_prob

    def calc_loss(self, rating_pred, rating_target, wishlist_pred, wishlist_target):
        mask_r = (wishlist_target == 0.0)
        rating_loss   = F.mse_loss(rating_pred[mask_r], rating_target[mask_r])
        wishlist_loss = F.mse_loss(wishlist_pred, wishlist_target)
        return rating_loss + self.wishlist_weight * wishlist_loss

def build_interaction_graph_from_dataset(
    dataset: torch.utils.data.TensorDataset,
    use_wishlist: bool = True
) -> torch.LongTensor:
    sids, pids, ratings, wishlists = dataset.tensors
    if use_wishlist:
        mask = wishlists == 1.0
    else:
        mask = torch.ones_like(wishlists, dtype=torch.bool)

    edge_index = torch.stack([sids[mask], pids[mask]], dim=0)
    return edge_index

train_dataset = get_dataset(train_ratings_df, train_wishlist_df)
valid_dataset = get_dataset(valid_ratings_df)

edge_index = build_interaction_graph_from_dataset(train_dataset, use_wishlist=False).to(device)

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
valid_loader = torch.utils.data.DataLoader(valid_dataset, batch_size=BATCH_SIZE, shuffle=False)

NUM_USERS       = 10_000
NUM_ITEMS       = 1_000
EMB_DIM         = 16
MLP_LAYERS      = [32, 32]
WISHLIST_WEIGHT = 2.0
BATCH_SIZE      = 64
LEARNING_RATE   = 1e-3

model_graphneumf = GraphNeuMF(
    num_users=NUM_USERS,
    num_items=NUM_ITEMS,
    emb_dim=EMB_DIM,
    mlp_layers=MLP_LAYERS,
    wishlist_weight=2.0
).to(device)

# precompute adjacency
model_graphneumf.lightgcn.set_adj(edge_index, device)

print("Pre-training LightGCN...")
pretrain_opt       = torch.optim.Adam(model_graphneumf.lightgcn.parameters(), lr=0.01)
GCN_PRETRAIN_EPOCHS= 5

for epoch in range(GCN_PRETRAIN_EPOCHS):
    model_graphneumf.lightgcn.train()
    pretrain_opt.zero_grad()

    user_emb, item_emb = model_graphneumf.lightgcn()
    pos_u, pos_i       = edge_index
    neg_i              = torch.randint(0, NUM_ITEMS, (edge_index.size(1)*10,), device=device)

    pos_scores = (user_emb[pos_u] * item_emb[pos_i]).sum(dim=1)
    neg_scores = (user_emb[pos_u.repeat(10)] * item_emb[neg_i]).sum(dim=1)

    loss = -torch.log(torch.sigmoid(pos_scores.repeat(10) - neg_scores)).mean()
    loss.backward()
    pretrain_opt.step()

    print(f" Pre-train Epoch {epoch+1}/{GCN_PRETRAIN_EPOCHS}  Loss={loss.item():.4f}")
    if loss.item() < 0.5:
        break

optimizer = torch.optim.AdamW(model_graphneumf.parameters(), lr=LEARNING_RATE)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
    optimizer, mode='min', factor=0.5, patience=1
)

NUM_EPOCHS = 40
for epoch in range(1, NUM_EPOCHS+1):
    model_graphneumf.train()
    tot_loss, tot_n = 0.0, 0

    for sid, pid, ratings, wishlists in train_loader:
        sid, pid = sid.to(device), pid.to(device)
        ratings, wishlists = ratings.to(device), wishlists.to(device)

        pr, pw = model_graphneumf(sid, pid)
        loss   = model_graphneumf.calc_loss(pr, ratings, pw, wishlists)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        tot_loss += loss.item() * sid.size(0)
        tot_n    += sid.size(0)

    avg_train_loss = tot_loss / tot_n

    # validation
    model_graphneumf.eval()
    val_mse, val_n = 0.0, 0
    with torch.no_grad():
        for sid, pid, ratings, _ in valid_loader:
            sid, pid   = sid.to(device), pid.to(device)
            ratings    = ratings.to(device)
            pred_r, _  = model_graphneumf(sid, pid)
            pred_r_cl  = pred_r.clamp(1.0, 5.0)
            val_mse   += F.mse_loss(pred_r_cl, ratings, reduction='sum').item()
            val_n     += sid.size(0)

    valid_rmse = (val_mse / val_n) ** 0.5
    print(f"[Epoch {epoch}/{NUM_EPOCHS}] "
          f"Train loss={avg_train_loss:.4f}, Valid RMSE={valid_rmse:.4f}")

    scheduler.step(valid_rmse)

pred_fn_graphneumf = lambda sids, pids: model_graphneumf(
    torch.from_numpy(sids).to(device),
    torch.from_numpy(pids).to(device)
)[0].clamp(1, 5).detach().cpu().numpy()

# Evaluate on validation data
with torch.no_grad():
    val_score = evaluate(valid_ratings_df, pred_fn_graphneumf)

print(f"Validation RMSE: {val_score:.3f}")

def make_block(in_f, out_f):
    return nn.Sequential(
        nn.Linear(in_f, out_f),
        nn.GELU(),
    )

class DMF(nn.Module):
    def __init__(
        self,
        num_users,
        num_items,
        layers,
        wishlist_weight=0.5,
        max_rating=5.0,
        embed_lambda=1e-5,
    ):
        super().__init__()
        self.max_rating           = max_rating
        self.wishlist_weight = wishlist_weight
        self.embed_lambda         = embed_lambda

        emb_dim = layers[0]
        self.u_emb       = nn.Embedding(num_users, emb_dim)
        self.i_emb       = nn.Embedding(num_items, emb_dim)
        self.u_bias      = nn.Embedding(num_users, 1)
        self.i_bias      = nn.Embedding(num_items, 1)
        self.global_bias = nn.Parameter(torch.zeros(1))

        self.dropout = nn.Dropout(0.2)

        self.user_tower = nn.Sequential(*[
            make_block(layers[i], layers[i+1])
            for i in range(len(layers) - 1)
        ])
        self.item_tower = nn.Sequential(*[
            make_block(layers[i], layers[i+1])
            for i in range(len(layers) - 1)
        ])

        self.fc_wish = nn.Linear(layers[-1], 1)

    @staticmethod
    def bpr_loss(pos_score, neg_score):
        return -F.logsigmoid(pos_score - neg_score).mean()

    def forward(self, u_idx, i_idx):
        u = self.dropout(self.u_emb(u_idx))
        v = self.dropout(self.i_emb(i_idx))

        u = self.user_tower(u)
        v = self.item_tower(v)

        x = u * v

        dot    = x.sum(dim=1)
        ub     = self.u_bias(u_idx).squeeze(1)
        ib     = self.i_bias(i_idx).squeeze(1)
        rating = dot + ub + ib + self.global_bias

        wish_logit = self.fc_wish(x).squeeze(1)
        wish       = torch.sigmoid(wish_logit)

        return rating, wish, wish_logit

    def calc_loss(
        self,
        rating_pred,
        rating_true,
        wish_logit_pos=None,
        wish_logit_neg=None,
    ):
        mask = rating_true > 0
        rating_loss = F.mse_loss(rating_pred[mask], rating_true[mask])

        if wish_logit_pos is None or wish_logit_neg is None:
            imp_loss = torch.tensor(0., device=rating_pred.device, dtype=rating_pred.dtype)
        else:
            imp_loss = self.bpr_loss(wish_logit_pos, wish_logit_neg)

        base = rating_loss + self.wishlist_weight * imp_loss

        emb_reg = (
            self.u_emb.weight.norm(p=2).pow(2)
          + self.i_emb.weight.norm(p=2).pow(2)
        ) * self.embed_lambda

        return base + 0.5 * emb_reg

train_dataset = get_dataset(train_ratings_df, train_wishlist_df)
valid_dataset = get_dataset(valid_ratings_df)

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)
valid_loader = torch.utils.data.DataLoader(valid_dataset, batch_size=64, shuffle=False)

NUM_USERS   = 10000
NUM_ITEMS   = 1000
MLP_HIDDEN  = [32, 64, 128]
LEARNING_RATE = 1e-3
WEIGHT_DECAY  = 1e-3
EMBED_LAMBDA  = 1e-4

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model_dmf = DMF(NUM_USERS, NUM_ITEMS, MLP_HIDDEN, wishlist_weight=0.5, max_rating=5.0, embed_lambda=EMBED_LAMBDA).to(device)
optimizer = torch.optim.AdamW(model_dmf.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=1)

NUM_EPOCHS = 40
for epoch in range(NUM_EPOCHS):
    model_dmf.train()
    running_loss = 0.0

    for sid, pid, ratings, wishlists in train_loader:
        sid, pid   = sid.to(device), pid.to(device)
        ratings    = ratings.to(device)
        wishlists  = wishlists.to(device)

        pred_ratings, _, wish_logit = model_dmf(sid, pid)

        # BPR negative sampling using wishlist flag
        pos_idx = (wishlists > 0).nonzero(as_tuple=True)[0]
        if pos_idx.numel() > 0:
            pos_scores = wish_logit[pos_idx]
            u_pos      = sid[pos_idx]
            neg_items  = torch.randint(0, NUM_ITEMS, (pos_idx.size(0),), device=device)
            _, _, neg_logits = model_dmf(u_pos, neg_items)
        else:
            pos_scores, neg_logits = None, None

        loss = model_dmf.calc_loss(
            pred_ratings, ratings,
            pos_scores, neg_logits
        )

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        running_loss += loss.item() * sid.size(0)

    avg_train_loss = running_loss / len(train_loader.dataset)

    model_dmf.eval()
    mse_sum = 0.0
    with torch.no_grad():
        for sid, pid, ratings, _ in valid_loader:
            sid, pid   = sid.to(device), pid.to(device)
            ratings    = ratings.to(device)
            pred_ratings, _, _ = model_dmf(sid, pid)
            pred_ratings       = pred_ratings.clamp(1, model_dmf.max_rating)
            mse_sum += F.mse_loss(pred_ratings, ratings, reduction='sum').item()

    rmse = (mse_sum / len(valid_loader.dataset)) ** 0.5
    print(f"Epoch {epoch+1}/{NUM_EPOCHS} – Train Loss: {avg_train_loss:.4f}, Valid RMSE: {rmse:.4f}")
    scheduler.step(rmse)

pred_fn_dmf = lambda sids, pids: (
    model_dmf(
        torch.from_numpy(sids).long().to(device),
        torch.from_numpy(pids).long().to(device)
    )[0].detach().cpu().numpy()
)

# Evaluate on validation data
with torch.no_grad():
    val_score = evaluate(valid_ratings_df, pred_fn_dmf)

print(f"Validation RMSE: {val_score:.3f}")

prediction_functions = {
    "NeuMF": pred_fn_neumf,
    "DMF": pred_fn_dmf,
    "SVD++": pred_fn_svdpp,
    "GraphNeuMF": pred_fn_graphneumf,
}

true_ratings_train = train_ratings_df["rating"].values
true_ratings_valid = valid_ratings_df["rating"].values

def get_model_predictions(df, model_name):
    sids, pids = df["sid"].values, df["pid"].values
    return prediction_functions[model_name](sids, pids)

train_predictions = {name: get_model_predictions(train_ratings_df, name) for name in prediction_functions}
valid_predictions = {name: get_model_predictions(valid_ratings_df, name) for name in prediction_functions}

base_model_rmses = {
    name: np.sqrt(np.mean((valid_predictions[name] - true_ratings_valid) ** 2))
    for name in prediction_functions
}

for name, rmse in base_model_rmses.items():
    print(f"{name} Validation RMSE: {rmse:.4f}")

best_base_model = min(base_model_rmses, key=base_model_rmses.get)
best_base_rmse = base_model_rmses[best_base_model]

# Simple Average
avg_preds = np.mean(list(valid_predictions.values()), axis=0)
avg_rmse = np.sqrt(np.mean((avg_preds - true_ratings_valid) ** 2))
print(f"Simple Average RMSE: {avg_rmse:.4f}")

# Weighted Average
inv_rmses = {k: 1/v for k, v in base_model_rmses.items()}
norm = sum(inv_rmses.values())
weights = {k: v / norm for k, v in inv_rmses.items()}

weighted_preds = sum(weights[k] * valid_predictions[k] for k in weights)
weighted_rmse = np.sqrt(np.mean((weighted_preds - true_ratings_valid) ** 2))
print(f"Weighted Average RMSE: {weighted_rmse:.4f}")

from sklearn.preprocessing import StandardScaler

def create_meta_features(pred_dict):
    preds = list(pred_dict.values())
    return np.vstack([
        preds[0],
        preds[1],
        np.mean(preds, axis=0),
        preds[0] - preds[1]
    ]).T

X_train = create_meta_features(train_predictions)
X_valid = create_meta_features(valid_predictions)

scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_valid_scaled = scaler.transform(X_valid)

X_train_tensor = torch.tensor(X_train_scaled, dtype=torch.float32).to(device)
X_valid_tensor = torch.tensor(X_valid_scaled, dtype=torch.float32).to(device)
y_train_tensor = torch.tensor(true_ratings_train, dtype=torch.float32).to(device)
y_valid_tensor = torch.tensor(true_ratings_valid, dtype=torch.float32).to(device)

class DeepStacking(nn.Module):
    def __init__(self, in_dim, p_drop=0.2):
        super().__init__()
        h1 = 4 * in_dim
        self.net = nn.Sequential(
            nn.Linear(in_dim, h1),
            nn.BatchNorm1d(h1),
            nn.GELU(),
            nn.Dropout(p_drop),
            nn.Linear(h1, 1),
        )

    def forward(self, x):
        return self.net(x).squeeze(-1)


# model, optimizer, loss
stack_model = DeepStacking(X_train.shape[1]).to(device)
optimizer = torch.optim.AdamW(stack_model.parameters(), lr=1e-3, weight_decay=1e-3)
loss_fn = nn.MSELoss()

# scheduler: reduce LR on plateau of val RMSE
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
    optimizer,
    mode='min',           # we want to minimize RMSE
    factor=0.5,           # LR = LR * factor
    patience=1 ,          # wait 10 epochs before reducing
)

best_rmse = float("inf")

for epoch in range(10000):
    # ——— Training step ———
    stack_model.train()
    optimizer.zero_grad()
    pred = stack_model(X_train_tensor)
    loss = loss_fn(pred, y_train_tensor)
    loss.backward()
    optimizer.step()

    # ——— Every 250 epochs: evaluate both train and val ———
    if epoch % 250 == 0:
        stack_model.eval()
        with torch.no_grad():
            # training RMSE
            train_preds = stack_model(X_train_tensor)
            train_rmse  = torch.sqrt(loss_fn(train_preds, y_train_tensor)).item()

            # validation RMSE (with clamp if desired)
            val_preds = stack_model(X_valid_tensor).clamp(1, 5)
            val_rmse  = torch.sqrt(loss_fn(val_preds, y_valid_tensor)).item()

            print(
                f"Epoch {epoch:4d} | "
                f"Train RMSE: {train_rmse:.4f} | "
                f"Val   RMSE: {val_rmse:.4f}"
            )

            scheduler.step(val_rmse)

            # checkpoint best
            if val_rmse < best_rmse:
                best_rmse = val_rmse
                best_state = stack_model.state_dict()

stack_model.load_state_dict(best_state)
stack_model.eval()
with torch.no_grad():
    final_preds = stack_model(X_valid_tensor).clamp(1, 5).cpu().numpy()
    final_rmse = np.sqrt(np.mean((final_preds - true_ratings_valid) ** 2))

print(f"\nFinal Linear Stacking RMSE: {final_rmse:.4f}")

def top_k_ensemble(k=3):

    sorted_models = sorted(base_model_rmses.items(), key=lambda x: x[1])
    top_models = sorted_models[:k]
    top_model_names = [name for name, _ in top_models]

    print(f"Selected top {k} models: {', '.join(top_model_names)}")


    weights = {}
    sum_weights = sum(1/ rmse for _, rmse in top_models)
    for name, rmse in top_models:
        weights[name] = (1 / rmse) / sum_weights
        print(f"  {name} weight: {weights[name]:.4f}")

    preds = np.zeros_like(true_ratings_valid, dtype=float)
    for name in top_model_names:
        preds += weights[name] * valid_predictions[name]

    rmse = np.sqrt(np.mean((preds - true_ratings_valid) ** 2))
    return preds, rmse, weights, top_model_names

best_k = 0
best_rmse = float('inf')

for k in range(2, len(prediction_functions) + 1):
    _, rmse, _, _ = top_k_ensemble(k)
    if rmse < best_rmse:
        best_rmse = rmse
        best_k = k

print(f"Best Top-{best_k} Ensemble RMSE: {best_rmse:.6f}")