"""Final version: LightGCN-style + learnable embeddings + full data training.

Key improvements over V2:
1. LightGCN-style: no ReLU, no feature transform between GNN layers (just aggregate)
2. Learnable author embeddings (replacing 2-dim degree features as primary input)
3. Layer combination (weighted sum of all layer outputs)
4. Multi-neg: 2 negatives per positive for better BPR signal
5. Train on full data for final submission
6. 5-model ensemble
7. Edge dropout for regularization
"""
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_geometric.data import HeteroData
from sklearn.metrics import 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)


# ── Load data ─────────────────────────────────────────────────────
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)

# Known positives
train_set = set(map(tuple, existing_refs))
overlap = train_set & set(map(tuple, refs_to_pred))
print(f"Known positives in test: {len(overlap)}")

# ── 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")

# ── Degree features (for augmentation only) ───────────────────────
author_ref_deg = np.zeros(num_authors, dtype=np.float32)
paper_ref_deg = np.zeros(num_papers, dtype=np.float32)
paper_cite_in = np.zeros(num_papers, dtype=np.float32)
paper_cite_out = 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 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)


paper_feat_np = paper_feature.numpy().astype(np.float32)
paper_deg_feat = np.stack([
    log_norm(paper_ref_deg),
    log_norm(paper_cite_out),
    log_norm(paper_cite_in),
], axis=-1)
paper_feat_aug = np.concatenate([paper_feat_np, paper_deg_feat], axis=-1)

# Normalize paper features for stability
paper_mean = paper_feat_aug.mean(axis=0)
paper_std = paper_feat_aug.std(axis=0) + 1e-8
paper_feat_aug = (paper_feat_aug - paper_mean) / paper_std

print(f"Paper features: {paper_feat_aug.shape[1]}d (normalized)")

# ── Popular papers for hard negatives ─────────────────────────────
popular_threshold = np.percentile(paper_ref_deg[paper_ref_deg > 0], 70)
popular_papers = np.where(paper_ref_deg >= popular_threshold)[0]

# 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_90 = ref_edges_idx.sample(frac=0.9, random_state=0, axis=0)
val_pos = ref_edges_idx[~ref_edges_idx.index.isin(train_refs_90.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)

# ── Build HeteroData ──────────────────────────────────────────────
def build_data(ref_edges_use, num_a, num_p, paper_feat, device):
    """Build HeteroData with given reference edges."""
    ref_tensor = torch.as_tensor(
        ref_edges_use[['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, dtype=torch.float)

    d = HeteroData()
    d['author'].num_nodes = num_a
    d['paper'].num_nodes = num_p
    d['paper'].x = paper_x
    d['author', 'ref', 'paper'].edge_index = ref_tensor.t().contiguous()
    d['paper', 'beref', 'author'].edge_index = ref_tensor[:, [1, 0]].t().contiguous()
    d['paper', 'cite', 'paper'].edge_index = torch.cat([
        cite_tensor, cite_tensor[:, [1, 0]],
    ], dim=0).t().contiguous()
    d['author', 'coauthor', 'author'].edge_index = torch.cat([
        coauthor_tensor, coauthor_tensor[:, [1, 0]],
    ], dim=0).t().contiguous()
    return d.to(device)


# ── LightGCN-style Model ──────────────────────────────────────────
class LightGCNLayer(nn.Module):
    """LightGCN propagation: no feature transform, no activation.
    Just normalized neighborhood mean aggregation."""
    def __init__(self, metadata):
        super().__init__()
        node_types, edge_types = metadata
        self.edge_types_used = [
            ('author', 'ref', 'paper'),
            ('paper', 'beref', 'author'),
            ('paper', 'cite', 'paper'),
            ('author', 'coauthor', 'author'),
        ]

    def forward(self, x_dict, edge_index_dict):
        out_dict = {}
        # Aggregate messages for each edge type
        agg_dict = {nt: [] for nt in x_dict}

        for et in self.edge_types_used:
            if et not in edge_index_dict:
                continue
            src_type, _, dst_type = et
            src, dst = edge_index_dict[et]
            src_x = x_dict[src_type]

            # Mean aggregation
            agg = src_x.new_zeros((x_dict[dst_type].size(0), src_x.size(-1)))
            deg = src_x.new_zeros((x_dict[dst_type].size(0), 1))
            agg.index_add_(0, dst, src_x[src])
            deg.index_add_(0, dst, torch.ones(
                (dst.numel(), 1), dtype=src_x.dtype, device=src_x.device))
            agg = agg / deg.clamp(min=1.0)
            agg_dict[dst_type].append(agg)

        for nt in x_dict:
            if agg_dict[nt]:
                out_dict[nt] = sum(agg_dict[nt]) / len(agg_dict[nt])
            else:
                out_dict[nt] = x_dict[nt]
        return out_dict


class LightGCNRecommender(nn.Module):
    """LightGCN-style heterogeneous graph recommender."""
    def __init__(self, metadata, paper_in_dim, embed_dim=256, num_layers=4,
                 author_embed_dim=256):
        super().__init__()
        self.author_emb = nn.Embedding(num_authors, author_embed_dim)
        self.paper_proj = nn.Linear(paper_in_dim, embed_dim)

        self.layers = nn.ModuleList(
            [LightGCNLayer(metadata) for _ in range(num_layers)])
        self.num_layers = num_layers
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.xavier_uniform_(self.author_emb.weight)
        nn.init.xavier_uniform_(self.paper_proj.weight)
        nn.init.zeros_(self.paper_proj.bias)

    def encode(self, data):
        x_dict = {
            'author': self.author_emb.weight,
            'paper': self.paper_proj(data['paper'].x),
        }
        # Store all layer embeddings for combination
        all_layers = [x_dict]
        for layer in self.layers:
            x_dict = layer(x_dict, data.edge_index_dict)
            all_layers.append(x_dict)

        # Weighted sum: later layers get higher weight
        weights = torch.tensor(
            [1.0 / (self.num_layers + 1)] * (self.num_layers + 1),
            device=x_dict['author'].device,
        )
        final = {
            nt: sum(w * layer[nt] for w, layer in zip(weights, all_layers))
            for nt in x_dict
        }
        return final

    def decode(self, z_dict, edge_index):
        src, dst = edge_index
        return (z_dict['author'][src] * z_dict['paper'][dst]).sum(dim=-1)


# ── Hard negative sampling ───────────────────────────────────────
def sample_hard_negatives(n_samples, existing_set, device):
    neg_list = []

    def add_random(target):
        nonlocal neg_list
        while len(neg_list) < target:
            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))

    # 50% random
    add_random(int(n_samples * 0.5))

    # 25% popular
    cnt = 0
    target_pop = int(n_samples * 0.75)
    while len(neg_list) < target_pop and cnt < n_samples * 2:
        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))

    # 25% co-author
    cnt = 0
    while len(neg_list) < n_samples and cnt < n_samples * 3:
        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(n_samples)
    return torch.tensor(neg_list[:n_samples], 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, val_df):
    model.eval()
    z_dict = model.encode(data)
    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_batched(model, data, test_pairs, batch_size=65536):
    model.eval()
    z_dict = model.encode(data)
    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)


# ── Training ──────────────────────────────────────────────────────
def train_model(seed, embed_dim=256, num_layers=4, lr=0.005,
                num_epochs=200, use_full_data=False):
    set_seed(seed)

    if use_full_data:
        train_edges = ref_edges  # full data
    else:
        train_edges = train_refs_90  # 90% for validation

    data_local = build_data(train_edges, num_authors, num_papers,
                            paper_feat_aug, device)

    model = LightGCNRecommender(
        data_local.metadata(),
        paper_in_dim=paper_feat_aug.shape[1],
        embed_dim=embed_dim,
        num_layers=num_layers,
    ).to(device)

    optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=1e-5)
    pos_edge_index = data_local['author', 'ref', 'paper'].edge_index
    existing_train = set(
        map(tuple, train_edges[['source', 'target']].to_numpy().tolist()))

    batch_size = min(32768, pos_edge_index.size(1))
    neg_per_pos = 2  # multiple negatives per positive
    best_val_f1 = 0
    best_state = None
    best_threshold = 0.5
    patience = 0

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

        perm = torch.randperm(pos_edge_index.size(1), device=device)[:batch_size]
        pos_batch = pos_edge_index[:, perm]
        neg_batch = sample_hard_negatives(
            pos_batch.size(1) * neg_per_pos, existing_train, device)

        z_dict = model.encode(data_local)
        pos_score = model.decode(z_dict, pos_batch)
        neg_score = model.decode(z_dict, neg_batch)

        # BPR loss with multiple negs
        pos_score = pos_score.repeat_interleave(neg_per_pos)
        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()

        if not use_full_data and (epoch % 10 == 0 or epoch == num_epochs - 1):
            val_f1, val_auc, val_thresh = evaluate(model, data_local, val_set)
            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 = 0
            else:
                patience += 1

            marker = '*' if val_f1 == best_val_f1 else ' '
            if epoch % 50 == 0 or val_f1 == best_val_f1:
                print(f'{marker} Seed={seed} Epoch {epoch:03d} | Loss={loss.item():.4f} | '
                      f'Val F1={val_f1:.4f} AUC={val_auc:.4f} Thresh={val_thresh:.4f}')

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

    if not use_full_data and best_state is not None:
        model.load_state_dict(best_state)

    return model, best_val_f1, best_threshold, data_local


# ── Phase 1: Train with validation to find best config ────────────
print("\n" + "=" * 60)
print("Phase 1: Training with validation (90% data)")
print("=" * 60)

models_val = []
thresholds_val = []
for seed in [0, 10, 42, 100, 2024]:
    print(f"\n--- Seed {seed} ---")
    m, f1, thresh, d = train_model(seed, embed_dim=256, num_layers=4,
                                   use_full_data=False)
    models_val.append(m)
    thresholds_val.append(thresh)
    print(f"  Best val F1: {f1:.4f}, Threshold: {thresh:.4f}")

avg_thresh = np.mean(thresholds_val)
print(f"\nAverage threshold: {avg_thresh:.4f}")

# ── Phase 2: Train on full data for final prediction ──────────────
print("\n" + "=" * 60)
print("Phase 2: Training on FULL data for submission")
print("=" * 60)

models_full = []
for seed in [0, 10, 42, 100, 2024]:
    print(f"\n--- Seed {seed} (full data) ---")
    m, _, _, d = train_model(seed, embed_dim=256, num_layers=4,
                             use_full_data=True, num_epochs=80)
    models_full.append(m)

# ── Generate submission ───────────────────────────────────────────
print("\n" + "=" * 60)
print("Generating final submission...")
print("=" * 60)

test_arr = np.array(refs_to_pred, dtype=np.int64)
data_full = build_data(ref_edges, num_authors, num_papers, paper_feat_aug, device)

all_scores = []
for i, model in enumerate(models_full):
    model = model.to(device)
    scores = predict_cos_batched(model, data_full, test_arr)
    all_scores.append(scores)
    print(f"  Model {i+1}: done, mean score={scores.mean():.4f}")

ensemble_scores = np.mean(all_scores, axis=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: use avg from validation phase
final_threshold = avg_thresh
print(f"Threshold: {final_threshold:.4f}")

predictions = (ensemble_scores >= final_threshold).astype(int)
print(f"Positive ratio: {predictions.mean():.4f} ({predictions.sum()}/{len(predictions)})")
print(f"Known positives set to 1: {known_pos_mask.sum()}")

# Save
output_path = "/home/lzc/submission_final.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"\nSaved to: {output_path}")