import pprint
from typing import List, Dict, Tuple
from typing import Optional, Any

from pydantic import BaseModel, Field, model_validator
from skellytracker.trackers.mediapipe_tracker.mediapipe_model_info import (
    MediapipeModelInfo,
)


class Point(BaseModel):
    """
    A point in 3D space at a particular time
    """

    x: Optional[float] = Field(None, description="The X-coordinate of the point")
    y: Optional[float] = Field(None, description="The Y-coordinate of the point")
    z: Optional[float] = Field(None, description="The Z-coordinate of the point")


class ReprojectionError(BaseModel):
    """
    Reprojection error of a 3D point at a particular time
    """

    value: Optional[float] = Field(None, description="Reprojection error of the point")


class VirtualMarkerDefinition(BaseModel):
    """
    A virtual marker, defined by combining multiple markers with weights to generate a new marker/point
    """

    marker_names: List[str] = Field(
        default_factory=list, description="The names of the markers that define this virtual marker"
    )
    marker_weights: List[float] = Field(
        default_factory=list, description="The weights of the markers that define this virtual marker, must sum to 1"
    )

    @model_validator(mode="before")
    def check_weights(cls, values):
        marker_weights = values.get("marker_weights")
        if sum(marker_weights) != 1:
            raise ValueError(f"Marker weights must sum to 1, got {marker_weights}")
        return values


class SegmentSchema(BaseModel):
    """
    A schema for a segment of a skeleton, defined by a set of tracked points and connections between them
    """

    point_names: List[str] = Field(
        default_factory=list, description="The names of the tracked points that define this segment"
    )
    virtual_marker_definitions: Optional[Dict[str, VirtualMarkerDefinition]] = Field(
        default_factory=dict, description="The virtual markers that define this segment"
    )
    connections: List[Tuple[int, int]] = Field(
        default_factory=list, description="The connections between the tracked points"
    )
    parent: Optional[str] = Field(None, description="The name of the parent of this segment")


class SkeletonSchema(BaseModel):
    body: SegmentSchema = Field(default_factory=SegmentSchema, description="The tracked points that define the body")
    hands: Dict[str, SegmentSchema] = Field(
        default_factory=dict, description="The tracked points that define the hands: keys - (left, right)"
    )
    face: SegmentSchema = Field(default_factory=SegmentSchema, description="The tracked points that define the face")

    def __init__(self, schema_dict: Dict[str, Dict[str, Any]]):
        super().__init__()
        self.body = SegmentSchema(**schema_dict["body"])
        self.hands = {hand: SegmentSchema(**hand_schema) for hand, hand_schema in schema_dict["hands"].items()}
        self.face = SegmentSchema(**schema_dict["face"])

    def dict(self):
        d = {}
        d["body"] = self.body.model_dump()
        d["hands"] = {hand: hand_schema.model_dump() for hand, hand_schema in self.hands.items()}
        d["face"] = self.face.model_dump()
        return d


class Timestamps(BaseModel):
    # TODO - Add Unix time and ISO format stuff and whatnot
    mean: Optional[float] = Field(None, description="The mean timestamp for this frame")
    by_camera: Dict[str, Any] = Field(
        default_factory=dict, description="Timestamps for each camera on this frame (key is video name)"
    )


class FrameData(BaseModel):
    """
    The data for a single frame
    """

    timestamps: Timestamps = Field(default_factory=Timestamps, description="Timestamp data")
    tracked_points: Dict[str, Point] = Field(default_factory=dict, description="The points being tracked")
    reprojection_error: Dict[str, ReprojectionError] = Field(
        default_factory=dict, description="Reprojection error of the point being tracked"
    )

    @property
    def tracked_point_names(self):
        return list(self.tracked_points.keys())

    @property
    def timestamp(self):
        return self.timestamps.mean

    def to_dict(self):
        return {
            "timestamps": self.timestamps.model_dump(),
            "tracked_points": {name: point.model_dump() for name, point in self.tracked_points.items()},
            "reprojection_error": {name: error.model_dump() for name, error in self.reprojection_error.items()},
        }


class InfoDict(BaseModel):
    """
    A dictionary of information about this recording, such as the measured segement lengths and the segment connections that we can use to interpret the tracked points (i.e./e.g. how to connect the dots of skeleton)
    """

    # segment_lengths: Dict[str, Any] = Field(default_factory=dict, description="The lengths of the segments of the body")
    schemas: Optional[Dict[str, Any]] = Field(
        default_factory=dict, description="The segment connections for the tracked points"
    )


if __name__ == "__main__":
    # nested_dict = create_nested_dict(FramePacket)
    # print(json.dumps(nested_dict, indent=4))

    print("=====================================\n\n===============================")

    skeleton_schema = SkeletonSchema(schema_dict=MediapipeModelInfo.skeleton_schema)

    pprint.pp(skeleton_schema.dict())
