from typing import Dict, List
import pandas as pd
# from .paper2graph import Graph
from .paper2graph import Graph

def create_entities(raw) -> pd.DataFrame:
    entities = []
    for unit in raw:
        e_unit = {
            "entity name": unit['title'],
            "entity type": "paper",
            "timetamp": unit['date'],
            "description": unit['abstract']
        }
        entities.append(e_unit)
    return pd.DataFrame(entities)

def create_relations(raw) -> pd.DataFrame:
    relations = []
    for unit in raw:
        for kw in unit['keywords']:
            r_unit = {
                "entity1": unit['title'],
                "relation": "related to",
                "entity2": kw
            }
            relations.append(r_unit)
    return pd.DataFrame(relations)

def connect(init_graph: Graph, new_info: List):
    '''
    The unit format of new_info:
        unit = {
            "title": title,
            "date": the publication year of the paper,
            "isAPA": True if the citations of the paper is in APA format
            "abstract": abstract,
            "related work": related work content (string type),
            "reference": a list including all the references,
            "keywords": a list including keywords (i.e. topic or problems)
        }
    '''
    _entity = init_graph.entity
    _relation = init_graph.relation

    new_entity = create_entities(new_info)
    new_relation = create_relations(new_info)

    entity = pd.concat([_entity, new_entity], ignore_index=True)
    relation = pd.concat([_relation, new_relation], ignore_index=True)

    return entity, relation


