"""Pydantic models for GA4GH Pedigree.

Models for representing family pedigree structures using
standardized relationship types from the KIN ontology.
"""

from __future__ import annotations

from enum import Enum

from ai4gh.core.types import GA4GHBase


class Sex(str, Enum):
    """Biological sex."""

    MALE = "MALE"
    FEMALE = "FEMALE"
    OTHER = "OTHER_SEX"
    UNKNOWN = "UNKNOWN_SEX"


class AffectedStatus(str, Enum):
    """Disease affected status."""

    AFFECTED = "AFFECTED"
    UNAFFECTED = "UNAFFECTED"
    MISSING = "MISSING"


class LifeStatus(str, Enum):
    """Individual life status."""

    ALIVE = "ALIVE"
    DECEASED = "DECEASED"
    UNKNOWN = "UNKNOWN_STATUS"


class KINRelationship(str, Enum):
    """KIN ontology relationship types."""

    PARENT = "KIN:001"
    MOTHER = "KIN:002"
    FATHER = "KIN:003"
    SIBLING = "KIN:004"
    TWIN = "KIN:005"
    MONOZYGOTIC_TWIN = "KIN:006"
    DIZYGOTIC_TWIN = "KIN:007"
    CHILD = "KIN:008"
    HALF_SIBLING = "KIN:009"
    ADOPTED_CHILD = "KIN:010"
    FOSTER_CHILD = "KIN:011"
    GRANDPARENT = "KIN:012"
    GRANDCHILD = "KIN:013"
    SPOUSE = "KIN:014"
    COUSIN = "KIN:015"
    UNCLE_AUNT = "KIN:016"
    NEPHEW_NIECE = "KIN:017"


class PedigreeIndividual(GA4GHBase):
    """An individual in a pedigree."""

    individual_id: str
    sex: Sex | None = None
    affected_status: AffectedStatus | None = None
    life_status: LifeStatus | None = None
    maternal_id: str | None = None
    paternal_id: str | None = None
    date_of_birth: str | None = None


class PedigreeRelationship(GA4GHBase):
    """A relationship between two individuals."""

    subject_id: str
    object_id: str
    relationship: KINRelationship


class Pedigree(GA4GHBase):
    """A family pedigree."""

    id: str | None = None
    proband_id: str | None = None
    individuals: list[PedigreeIndividual]
    relationships: list[PedigreeRelationship] | None = None

    def get_individual(self, individual_id: str) -> PedigreeIndividual | None:
        """Find an individual by ID."""
        for ind in self.individuals:
            if ind.individual_id == individual_id:
                return ind
        return None

    def get_parents(self, individual_id: str) -> list[PedigreeIndividual]:
        """Get parents of an individual."""
        ind = self.get_individual(individual_id)
        if ind is None:
            return []
        parents = []
        if ind.maternal_id:
            mom = self.get_individual(ind.maternal_id)
            if mom:
                parents.append(mom)
        if ind.paternal_id:
            dad = self.get_individual(ind.paternal_id)
            if dad:
                parents.append(dad)
        return parents

    def get_siblings(self, individual_id: str) -> list[PedigreeIndividual]:
        """Get siblings of an individual (shared at least one parent)."""
        ind = self.get_individual(individual_id)
        if ind is None:
            return []
        siblings = []
        for other in self.individuals:
            if other.individual_id == individual_id:
                continue
            if (ind.maternal_id and other.maternal_id == ind.maternal_id) or (
                ind.paternal_id and other.paternal_id == ind.paternal_id
            ):
                siblings.append(other)
        return siblings

    def get_affected(self) -> list[PedigreeIndividual]:
        """Get all affected individuals."""
        return [ind for ind in self.individuals if ind.affected_status == AffectedStatus.AFFECTED]
