import random

import torch
import dgl
import pandas as pd
import numpy as np
from tqdm.auto import tqdm

from exp.utils import normalize_embeddings


def fix_random(seed):
    dgl.seed(seed)
    torch.random.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)


class LRSchedule:
    def __init__(self, total_steps, warmup_steps, final_factor):
        self._total_steps = total_steps
        self._warmup_steps = warmup_steps
        self._final_factor = final_factor
        
    def __call__(self, step):
        if step >= self._total_steps:
            return self._final_factor
        
        if self._warmup_steps > 0:
            warmup_factor = step / self._warmup_steps
        else:
            warmup_factor = 1.0
        
        steps_after_warmup = step - self._warmup_steps
        total_steps_after_warmup = self._total_steps - self._warmup_steps
        after_warmup_factor = 1 \
            - (1 - self._final_factor) * (steps_after_warmup / total_steps_after_warmup)
        
        factor = min(warmup_factor, after_warmup_factor)
        return min(max(factor, 0), 1)


def prepare_graphs(items_path, ratings_path):
    items = pd.read_csv(items_path)
    ratings = pd.read_csv(ratings_path)

    n_users = np.max(ratings["user_id"].unique()) + 1
    item_ids = torch.tensor(sorted(items["item_id"].unique()))

    edges = torch.tensor(ratings["user_id"]), torch.tensor(ratings["item_id"])
    reverse_edges = (edges[1], edges[0])

    bipartite_graph = dgl.heterograph(
        data_dict={
            ("User", "UserItem", "Item"): edges,
            ("Item", "ItemUser", "User"): reverse_edges
        },
        num_nodes_dict={
            "User": n_users,
            "Item": len(item_ids)
        }
    )
    graph = dgl.to_homogeneous(bipartite_graph)
    graph = dgl.add_self_loop(graph)
    return bipartite_graph, graph


def sample_item_batch(user_batch, bipartite_graph):
    sampled_edges = dgl.sampling.sample_neighbors(
        bipartite_graph, {"User": user_batch}, fanout=2
    ).edges(etype="ItemUser")
    item_batch = sampled_edges[0]
    item_batch = item_batch[torch.argsort(sampled_edges[1])]
    item_batch = item_batch.reshape(-1, 2)
    item_batch = item_batch.T
    return item_batch


@torch.no_grad()
def inference_model(model, bipartite_graph, batch_size, hidden_dim, device):
    model.eval()
    item_embeddings = torch.zeros(bipartite_graph.num_nodes("Item"), hidden_dim).to(device)
    for items_batch in tqdm(torch.utils.data.DataLoader(
            torch.arange(bipartite_graph.num_nodes("Item")), 
            batch_size=batch_size, 
            shuffle=True
    )):
        item_embeddings[items_batch] = model(items_batch.to(device))

    item_embeddings = normalize_embeddings(item_embeddings.cpu().numpy())
    return item_embeddings