"""V2: Strong encoder + dot product decoder + BPR loss + hard negatives.

Key insight from baseline: simple decoder (dot product) + normalized similarity
generalizes better. We improve the ENCODER and TRAINING instead.

Improvements over baseline:
1. SAGEConv-based heterogeneous GNN encoder (3 layers, residual, LayerNorm)
2. BPR loss (ranking-based, better for recommendation)
3. Hard negative sampling (popular + co-author papers)
4. Graph structural features augmented to paper features
5. Cosine similarity for prediction (normalized dot product)
6. Known positives from train set forced to 1
7. Ensemble multiple seeds
8. Longer training with early stopping
"""
import os
import pickle as pkl
import random

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.optim.lr_scheduler import ReduceLROnPlateau

from torch_geometric.data import HeteroData
from torch_geometric.nn import SAGEConv, HeteroConv
from sklearn.metrics import f1_score, precision_recall_curve, roc_auc_score
from numpy.linalg import norm

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print('device:', device)


def set_seed(seed=0):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


# ── Data loading ──────────────────────────────────────────────────
base_path = "/home/lzc/cs3319-project"


def read_txt(file):
    res_list = []
    with open(file, "r") as f:
        for line in f:
            res_list.append(list(map(int, line.strip().split())))
    return res_list


citation = read_txt(os.path.join(base_path, "paper_file_ann.txt"))
existing_refs = read_txt(os.path.join(base_path, "bipartite_train_ann.txt"))
refs_to_pred = read_txt(os.path.join(base_path, "bipartite_test_ann.txt"))
coauthor = read_txt(os.path.join(base_path, "author_file_ann.txt"))

with open(os.path.join(base_path, "feature.pkl"), 'rb') as f:
    paper_feature = pkl.load(f)

print(f"Data: {len(citation)} cites, {len(existing_refs)} refs, "
      f"{len(refs_to_pred)} test, {len(coauthor)} coauthor")

# ── Train-test overlap ────────────────────────────────────────────
train_set = set(map(tuple, existing_refs))
test_arr_full = np.array(refs_to_pred, dtype=np.int64)
overlap = train_set & set(map(tuple, refs_to_pred))
print(f"Known positives in test: {len(overlap)} ({100*len(overlap)/len(refs_to_pred):.1f}%)")

# ── Degree features ───────────────────────────────────────────────
# Build node sets
cite_edges = pd.DataFrame(citation, columns=['source', 'target'])
ref_edges = pd.DataFrame(existing_refs, columns=['source', 'target'])
coauthor_edges = pd.DataFrame(coauthor, columns=['source', 'target'])

node_tmp = pd.concat([cite_edges['source'], cite_edges['target'], ref_edges['target']])
node_papers = pd.DataFrame(index=pd.unique(node_tmp))
node_tmp = pd.concat([ref_edges['source'], coauthor_edges['source'], coauthor_edges['target']])
node_authors = pd.DataFrame(index=pd.unique(node_tmp))

num_authors = len(node_authors)
num_papers = len(node_papers)
print(f"Nodes: {num_authors} authors, {num_papers} papers")

# Compute degree features
author_ref_deg = np.zeros(num_authors, dtype=np.float32)
paper_ref_deg = np.zeros(num_papers, dtype=np.float32)
author_coauthor_deg = np.zeros(num_authors, dtype=np.float32)
paper_cite_out = np.zeros(num_papers, dtype=np.float32)
paper_cite_in = np.zeros(num_papers, dtype=np.float32)

for s, t in existing_refs:
    author_ref_deg[s] += 1
    paper_ref_deg[t] += 1
for s, t in coauthor:
    author_coauthor_deg[s] += 1
    author_coauthor_deg[t] += 1
for s, t in citation:
    paper_cite_out[s] += 1
    paper_cite_in[t] += 1


def log_norm(x):
    x = np.log1p(x)
    return (x - x.mean()) / (x.std() + 1e-8)


author_deg_feat = np.stack([
    log_norm(author_ref_deg),
    log_norm(author_coauthor_deg),
], axis=-1)

paper_deg_feat = np.stack([
    log_norm(paper_ref_deg),
    log_norm(paper_cite_out),
    log_norm(paper_cite_in),
], axis=-1)

# Augment paper features
paper_feat_aug = np.concatenate([paper_feature.numpy(), paper_deg_feat], axis=-1)
print(f"Feature dims: paper={paper_feat_aug.shape[1]}, author={author_deg_feat.shape[1]}")

# ── Pre-compute hard negative pools ───────────────────────────────
# Popular papers
popular_threshold = np.percentile(paper_ref_deg[paper_ref_deg > 0], 70)
popular_papers = np.where(paper_ref_deg >= popular_threshold)[0]
print(f"Popular papers: {len(popular_papers)}")

# Co-author paper pools
coauthor_map = {i: set() for i in range(num_authors)}
for s, t in coauthor:
    coauthor_map[s].add(t)
    coauthor_map[t].add(s)

author_papers = {i: set() for i in range(num_authors)}
for s, t in existing_refs:
    author_papers[s].add(t)

coauthor_paper_pool = {}
for author in range(num_authors):
    pool = set()
    for coa in coauthor_map[author]:
        pool.update(author_papers[coa])
    pool -= author_papers[author]
    coauthor_paper_pool[author] = list(pool) if pool else list(range(num_papers))

# ── Train/val split ──────────────────────────────────────────────
ref_edges_idx = ref_edges.copy()
train_refs = ref_edges_idx.sample(frac=0.9, random_state=0, axis=0)
val_pos = ref_edges_idx[~ref_edges_idx.index.isin(train_refs.index)].copy()
val_pos['label'] = 1

existing_ref_set = set(map(tuple, existing_refs))
author_ids_arr = node_authors.index.to_numpy(dtype=np.int64)
paper_ids_arr = node_papers.index.to_numpy(dtype=np.int64)

neg_pairs = []
rng = np.random.default_rng(0)
while len(neg_pairs) < len(val_pos):
    src = int(rng.choice(author_ids_arr))
    dst = int(rng.choice(paper_ids_arr))
    if (src, dst) not in existing_ref_set:
        neg_pairs.append((src, dst))

val_neg = pd.DataFrame(neg_pairs, columns=['source', 'target'])
val_neg['label'] = 0
val_set = pd.concat([val_pos, val_neg], ignore_index=True).sample(frac=1, random_state=0)
print(f"Val: {len(val_set)} pairs ({val_set['label'].sum()} pos)")

# ── Build HeteroData ──────────────────────────────────────────────
train_ref_tensor = torch.as_tensor(train_refs[['source', 'target']].to_numpy(), dtype=torch.long)
cite_tensor = torch.as_tensor(cite_edges[['source', 'target']].to_numpy(), dtype=torch.long)
coauthor_tensor = torch.as_tensor(coauthor_edges[['source', 'target']].to_numpy(), dtype=torch.long)

paper_x = torch.as_tensor(paper_feat_aug, dtype=torch.float)
author_x = torch.as_tensor(author_deg_feat, dtype=torch.float)

data = HeteroData()
data['author'].num_nodes = num_authors
data['author'].x = author_x
data['paper'].num_nodes = num_papers
data['paper'].x = paper_x
data['author', 'ref', 'paper'].edge_index = train_ref_tensor.t().contiguous()
data['paper', 'beref', 'author'].edge_index = train_ref_tensor[:, [1, 0]].t().contiguous()
data['paper', 'cite', 'paper'].edge_index = torch.cat([
    cite_tensor, cite_tensor[:, [1, 0]],
], dim=0).t().contiguous()
data['author', 'coauthor', 'author'].edge_index = torch.cat([
    coauthor_tensor, coauthor_tensor[:, [1, 0]],
], dim=0).t().contiguous()
data = data.to(device)


# ── Model ─────────────────────────────────────────────────────────
class GNNEncoder(nn.Module):
    """SAGEConv-based heterogeneous GNN with residuals and LayerNorm."""
    def __init__(self, metadata, author_in_dim, paper_in_dim,
                 hidden_dim=128, num_layers=3, dropout=0.2):
        super().__init__()
        node_types, edge_types = metadata
        edge_types_used = [
            ('author', 'ref', 'paper'),
            ('paper', 'beref', 'author'),
            ('paper', 'cite', 'paper'),
            ('author', 'coauthor', 'author'),
        ]

        self.author_proj = nn.Linear(author_in_dim, hidden_dim)
        self.paper_proj = nn.Linear(paper_in_dim, hidden_dim)

        self.convs = nn.ModuleList()
        self.norms = nn.ModuleList()
        self.dropout = nn.Dropout(dropout)

        for _ in range(num_layers):
            conv_dict = {}
            for et in edge_types_used:
                if et in edge_types:
                    conv_dict[et] = SAGEConv(hidden_dim, hidden_dim)
            self.convs.append(HeteroConv(conv_dict, aggr='mean'))
            self.norms.append(nn.ModuleDict({
                'author': nn.LayerNorm(hidden_dim),
                'paper': nn.LayerNorm(hidden_dim),
            }))

    def forward(self, x_dict, edge_index_dict):
        x_dict = {
            'author': self.author_proj(x_dict['author']),
            'paper': self.paper_proj(x_dict['paper']),
        }
        for conv, norm in zip(self.convs, self.norms):
            h = conv(x_dict, edge_index_dict)
            x_dict = {
                nt: self.dropout(F.relu(norm[nt](h[nt] + x_dict[nt])))
                for nt in h
            }
        return x_dict


class DotDecoder(nn.Module):
    """Simple dot product decoder."""
    def forward(self, author_emb, paper_emb, edge_index):
        src, dst = edge_index
        return (author_emb[src] * paper_emb[dst]).sum(dim=-1)


class Recommender(nn.Module):
    def __init__(self, metadata, author_in_dim, paper_in_dim,
                 hidden_dim=128, num_layers=3):
        super().__init__()
        self.encoder = GNNEncoder(metadata, author_in_dim, paper_in_dim,
                                  hidden_dim, num_layers)
        self.decoder = DotDecoder()
        self.reset_parameters()

    def reset_parameters(self):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.xavier_uniform_(m.weight)
                if m.bias is not None:
                    nn.init.zeros_(m.bias)

    def encode(self, x_dict, edge_index_dict):
        return self.encoder(x_dict, edge_index_dict)

    def decode(self, z_dict, edge_index):
        return self.decoder(z_dict['author'], z_dict['paper'], edge_index)


# ── Hard negative sampling ───────────────────────────────────────
def sample_hard_negatives(n_samples, existing_set, device):
    """Mixed: 60% random, 20% popular, 20% co-author papers."""
    neg_list = []

    def add_random(n):
        nonlocal neg_list
        while len(neg_list) < n:
            s = np.random.randint(0, num_authors)
            d = np.random.randint(0, num_papers)
            if (s, d) not in existing_set:
                neg_list.append((s, d))

    add_random(int(n_samples * 0.6))

    # Popular
    cnt = 0
    target = n_samples
    while len(neg_list) < int(n_samples * 0.8) and cnt < n_samples * 3:
        cnt += 1
        s = np.random.randint(0, num_authors)
        d = popular_papers[np.random.randint(0, len(popular_papers))]
        if (s, d) not in existing_set:
            neg_list.append((s, d))

    # Co-author
    cnt = 0
    while len(neg_list) < target and cnt < n_samples * 5:
        cnt += 1
        s = np.random.randint(0, num_authors)
        pool = coauthor_paper_pool.get(s, [])
        if pool:
            d = pool[np.random.randint(0, len(pool))]
            if (s, d) not in existing_set:
                neg_list.append((s, d))

    add_random(target)
    return torch.tensor(neg_list[:target], dtype=torch.long,
                        device=device).t().contiguous()


# ── Evaluation ────────────────────────────────────────────────────
def cos_sim(a, b, eps=1e-12):
    return np.sum(a * b, axis=1) / (norm(a, axis=1) * norm(b, axis=1) + eps)


@torch.no_grad()
def evaluate(model, data_dict, val_df):
    model.eval()
    z_dict = model.encode(data_dict, data.edge_index_dict)
    z_cpu = {k: v.cpu() for k, v in z_dict.items()}

    val_arr = val_df[['source', 'target']].to_numpy(dtype=np.int64)
    scores = cos_sim(
        z_cpu['author'][val_arr[:, 0]].numpy(),
        z_cpu['paper'][val_arr[:, 1]].numpy(),
    )
    labels = val_df['label'].to_numpy()

    precision, recall, thresholds = precision_recall_curve(labels, scores)
    f1s = 2 * precision * recall / (precision + recall + 1e-12)
    best_idx = np.argmax(f1s)
    best_thresh = thresholds[best_idx] if best_idx < len(thresholds) else 0.5
    best_f1 = f1s[best_idx]
    auc = roc_auc_score(labels, scores)
    return best_f1, auc, best_thresh


@torch.no_grad()
def predict_cos(model, data_dict, test_pairs):
    model.eval()
    z_dict = model.encode(data_dict, data.edge_index_dict)
    z_cpu = {k: v.cpu() for k, v in z_dict.items()}
    return cos_sim(
        z_cpu['author'][test_pairs[:, 0]].numpy(),
        z_cpu['paper'][test_pairs[:, 1]].numpy(),
    )


# ── Training ──────────────────────────────────────────────────────
def run_experiment(seed, hidden_dim=128, num_layers=3, lr=0.005,
                   num_epochs=300, use_hard_neg=True):
    set_seed(seed)
    model = Recommender(
        data.metadata(),
        author_in_dim=author_deg_feat.shape[1],
        paper_in_dim=paper_feat_aug.shape[1],
        hidden_dim=hidden_dim,
        num_layers=num_layers,
    ).to(device)

    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4)
    scheduler = ReduceLROnPlateau(optimizer, mode='max', factor=0.5,
                                  patience=15, min_lr=1e-6)

    existing_train_set = set(
        map(tuple, train_refs[['source', 'target']].to_numpy().tolist()))
    pos_edge_index = data['author', 'ref', 'paper'].edge_index
    batch_size = min(32768, pos_edge_index.size(1))

    best_val_f1 = 0
    best_state = None
    patience_counter = 0

    x_dict = {'author': data['author'].x, 'paper': data['paper'].x}

    for epoch in range(num_epochs):
        model.train()

        # Sample positive batch
        perm = torch.randperm(pos_edge_index.size(1), device=device)[:batch_size]
        pos_batch = pos_edge_index[:, perm]

        # Sample negatives
        if use_hard_neg:
            neg_batch = sample_hard_negatives(
                pos_batch.size(1), existing_train_set, device)
        else:
            neg_list = []
            while len(neg_list) < pos_batch.size(1):
                s = np.random.randint(0, num_authors)
                d = np.random.randint(0, num_papers)
                if (s, d) not in existing_train_set:
                    neg_list.append((s, d))
            neg_batch = torch.tensor(neg_list, dtype=torch.long,
                                     device=device).t().contiguous()

        # BPR loss
        z_dict = model.encode(x_dict, data.edge_index_dict)
        pos_score = model.decode(z_dict, pos_batch)
        neg_score = model.decode(z_dict, neg_batch)
        loss = -F.logsigmoid(pos_score - neg_score).mean()

        optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()

        # Validation
        if epoch % 5 == 0 or epoch == num_epochs - 1:
            val_f1, val_auc, val_thresh = evaluate(model, x_dict, val_set)
            scheduler.step(val_f1)

            if val_f1 > best_val_f1:
                best_val_f1 = val_f1
                best_state = {k: v.cpu().clone() for k, v in model.state_dict().items()}
                best_threshold = val_thresh
                patience_counter = 0
            else:
                patience_counter += 1

            if epoch % 20 == 0 or epoch == num_epochs - 1 or val_f1 == best_val_f1:
                status = '*' if val_f1 == best_val_f1 else ' '
                print(f'{status} Epoch {epoch:03d} | Loss={loss.item():.4f} | '
                      f'Val F1={val_f1:.4f} AUC={val_auc:.4f} Thresh={val_thresh:.3f}')

            if patience_counter >= 30:
                print(f'Early stopping at epoch {epoch}')
                break

    model.load_state_dict(best_state)
    return model, best_val_f1, best_threshold


# ── Run experiments ───────────────────────────────────────────────
print("\n" + "=" * 60)
print("Training model 1 (seed=0)")
print("=" * 60)
model1, f1_1, thresh1 = run_experiment(seed=0, hidden_dim=128, num_layers=3)

print("\n" + "=" * 60)
print("Training model 2 (seed=42)")
print("=" * 60)
model2, f1_2, thresh2 = run_experiment(seed=42, hidden_dim=128, num_layers=3)

print(f"\nModel 1: val F1={f1_1:.4f} thresh={thresh1:.4f}")
print(f"Model 2: val F1={f1_2:.4f} thresh={thresh2:.4f}")

# ── Ensemble prediction ───────────────────────────────────────────
print("\n" + "=" * 60)
print("Generating ensemble submission...")
print("=" * 60)


@torch.no_grad()
def predict_batched(model, data_dict, test_pairs, batch_size=65536):
    model.eval()
    z_dict = model.encode(data_dict, data.edge_index_dict)
    z_cpu = {k: v.cpu() for k, v in z_dict.items()}
    all_scores = []
    for start in range(0, len(test_pairs), batch_size):
        end = min(start + batch_size, len(test_pairs))
        batch = test_pairs[start:end]
        scores = cos_sim(
            z_cpu['author'][batch[:, 0]].numpy(),
            z_cpu['paper'][batch[:, 1]].numpy(),
        )
        all_scores.append(scores)
    return np.concatenate(all_scores)


x_dict = {'author': data['author'].x, 'paper': data['paper'].x}
test_arr = np.array(refs_to_pred, dtype=np.int64)

scores1 = predict_batched(model1, x_dict, test_arr)
scores2 = predict_batched(model2, x_dict, test_arr)
ensemble_scores = (scores1 + scores2) / 2.0

# Force known positives to 1
known_pos_mask = np.array([tuple(p) in overlap for p in refs_to_pred])
ensemble_scores[known_pos_mask] = 1.0

# Threshold from validation ensemble
val_s1 = predict_batched(model1, x_dict,
                         val_set[['source', 'target']].to_numpy(dtype=np.int64))
val_s2 = predict_batched(model2, x_dict,
                         val_set[['source', 'target']].to_numpy(dtype=np.int64))
val_ens = (val_s1 + val_s2) / 2.0
val_labels = val_set['label'].to_numpy()

precision, recall, thresholds = precision_recall_curve(val_labels, val_ens)
f1s = 2 * precision * recall / (precision + recall + 1e-12)
best_idx = np.argmax(f1s)
best_thresh = thresholds[best_idx] if best_idx < len(thresholds) else 0.5
print(f"Ensemble val F1: {f1s[best_idx]:.4f} @ thresh={best_thresh:.4f}")

# Apply threshold
predictions = (ensemble_scores >= best_thresh).astype(int)
print(f"Positive ratio: {predictions.mean():.4f} ({predictions.sum()}/{len(predictions)})")
print(f"Known positives: {known_pos_mask.sum()}")

# Save
output_path = "/home/lzc/submission_v2.csv"
data_out = [[idx, str(int(p))] for idx, p in enumerate(predictions)]
df = pd.DataFrame(data_out, columns=['Index', 'Predicted'], dtype=object)
df.to_csv(output_path, index=False)
print(f"Saved to: {output_path}")

# Also save individual model submissions
for name, scores_i in [('v2_m1', scores1), ('v2_m2', scores2)]:
    s = scores_i.copy()
    s[known_pos_mask] = 1.0
    preds = (s >= best_thresh).astype(int)
    out_path = f"/home/lzc/submission_{name}.csv"
    data_out = [[idx, str(int(p))] for idx, p in enumerate(preds)]
    pd.DataFrame(data_out, columns=['Index', 'Predicted'], dtype=object).to_csv(
        out_path, index=False)
    print(f"  {name}: {out_path}")