from enum import StrEnum
from typing import Sequence

from ie_datasets.util.interfaces import ImmutableModel


class DocREDEntityTypeID(StrEnum):
    PERSON = "PER"
    ORGANIZATION = "ORG"
    LOCATION = "LOC"
    TIME = "TIME"
    NUMBER = "NUM"
    OTHER = "MISC"


class DocREDRelationTypeID(StrEnum):
    HEAD_OF_GOVERNMENT = "P6"
    COUNTRY = "P17"
    PLACE_OF_BIRTH = "P19"
    PLACE_OF_DEATH = "P20"
    FATHER = "P22"
    MOTHER = "P25"
    SPOUSE = "P26"
    COUNTRY_OF_CITIZENSHIP = "P27"
    CONTINENT = "P30"
    INSTANCE_OF = "P31"
    HEAD_OF_STATE = "P35"
    CAPITAL = "P36"
    OFFICIAL_LANGUAGE = "P37"
    POSITION_HELD = "P39"
    CHILD = "P40"
    AUTHOR = "P50"
    MEMBER_OF_SPORTS_TEAM = "P54"
    DIRECTOR = "P57"
    SCREENWRITER = "P58"
    EDUCATED_AT = "P69"
    COMPOSER = "P86"
    MEMBER_OF_POLITICAL_PARTY = "P102"
    EMPLOYER = "P108"
    FOUNDED_BY = "P112"
    LEAGUE = "P118"
    PUBLISHER = "P123"
    OWNED_BY = "P127"
    LOCATED_IN_THE_ADMINISTRATIVE_TERRITORIAL_ENTITY = "P131"
    GENRE = "P136"
    OPERATOR = "P137"
    RELIGION = "P140"
    CONTAINS_ADMINISTRATIVE_TERRITORIAL_ENTITY = "P150"
    FOLLOWS = "P155"
    FOLLOWED_BY = "P156"
    HEADQUARTERS_LOCATION = "P159"
    CAST_MEMBER = "P161"
    PRODUCER = "P162"
    AWARD_RECEIVED = "P166"
    CREATOR = "P170"
    PARENT_TAXON = "P171"
    ETHNIC_GROUP = "P172"
    PERFORMER = "P175"
    MANUFACTURER = "P176"
    DEVELOPER = "P178"
    SERIES = "P179"
    SISTER_CITY = "P190"
    LEGISLATIVE_BODY = "P194"
    BASIN_COUNTRY = "P205"
    LOCATED_IN_OR_NEXT_TO_BODY_OF_WATER = "P206"
    MILITARY_BRANCH = "P241"
    RECORD_LABEL = "P264"
    PRODUCTION_COMPANY = "P272"
    LOCATION = "P276"
    SUBCLASS_OF = "P279"
    SUBSIDIARY = "P355"
    PART_OF = "P361"
    ORIGINAL_LANGUAGE_OF_WORK = "P364"
    PLATFORM = "P400"
    MOUTH_OF_THE_WATERCOURSE = "P403"
    ORIGINAL_NETWORK = "P449"
    MEMBER_OF = "P463"
    CHAIRPERSON = "P488"
    COUNTRY_OF_ORIGIN = "P495"
    HAS_PART = "P527"
    RESIDENCE = "P551"
    DATE_OF_BIRTH = "P569"
    DATE_OF_DEATH = "P570"
    INCEPTION = "P571"
    DISSOLVED_ABOLISHED_OR_DEMOLISHED = "P576"
    PUBLICATION_DATE = "P577"
    START_TIME = "P580"
    END_TIME = "P582"
    POINT_IN_TIME = "P585"
    CONFLICT = "P607"
    CHARACTERS = "P674"
    LYRICS_BY = "P676"
    LOCATED_ON_TERRAIN_FEATURE = "P706"
    PARTICIPANT = "P710"
    INFLUENCED_BY = "P737"
    LOCATION_OF_FORMATION = "P740"
    PARENT_ORGANIZATION = "P749"
    NOTABLE_WORK = "P800"
    SEPARATED_FROM = "P807"
    NARRATIVE_LOCATION = "P840"
    WORK_LOCATION = "P937"
    APPLIES_TO_JURISDICTION = "P1001"
    PRODUCT_OR_MATERIAL_PRODUCED = "P1056"
    UNEMPLOYMENT_RATE = "P1198"
    TERRITORY_CLAIMED_BY = "P1336"
    PARTICIPANT_OF = "P1344"
    REPLACES = "P1365"
    REPLACED_BY = "P1366"
    CAPITAL_OF = "P1376"
    LANGUAGES_SPOKEN_WRITTEN_OR_SIGNED = "P1412"
    PRESENT_IN_WORK = "P1441"
    SIBLING = "P3373"


class DocREDRelationType(ImmutableModel):
    id: DocREDRelationTypeID
    description: str


class DocREDSchema(ImmutableModel):
    relation_types: Sequence[DocREDRelationType]


__all__ = [
    "DocREDEntityTypeID",
    "DocREDRelationType",
    "DocREDRelationTypeID",
    "DocREDSchema",
]
