import torch
from datetime import datetime
from typing import Counter


def train_seed_generator(graph, args):
    """Generate Seed (Edges) based on some techniques
 
    [1] chronologically split, use fresh data (e.g. latest one year)
    [2] filter by edge type, here use the edge features dim=0, may use edge type explicitly
    """
 
    if args.edge_select == 'full':
        # use full edges without any filter strategies
        train_seed = torch.arange(graph.num_edges())
 
    elif args.edge_select == 'chronological':
        # use latest data for training
        ts_cut = datetime.timestamp(datetime.strptime(args.split_date, "%Y-%m-%d"))
        train_seed = torch.nonzero(
            graph.edata['timestamp'] >= ts_cut, 
            as_tuple=False
        ).squeeze()
 
    elif args.edge_select == 'etype_filter':
        # filter edges by edge type
 
        # here use edge features instand, or use graph.edata['etype'] explicitly
        edge_type_tensor = graph.edata['feats'][:, -3].long()
 
        # get target edge types:
        # [1] hot edge types topK (count)
        etype_cnt = Counter(edge_type_tensor.numpy())
        etype_sorted = sorted(etype_cnt.items(), key=lambda x: x[1], reverse=True)
        etype_selected = [x[0] for x in etype_sorted[:args.etype_topk]]
 
        # [2] hot edge types (percentage)， 用 pandas 处理会更方便一些
 
        tmp = edge_type_tensor.unsqueeze(-1).repeat(1, len(etype_selected))
        train_seed = (tmp == torch.tensor(etype_selected)).sum(-1).nonzero(as_tuple=False).squeeze()
 
    else:
        raise Exception(f"invalid edge select option: '{args.edge_select}', not implemented!")
 
    return train_seed