# Copyright (c) 2025 CoReason, Inc.
#
# This software is proprietary and dual-licensed.
# Licensed under the Prosperity Public License 3.0 (the "License").
# A copy of the license is available at https://prosperitylicense.com/versions/3.0.0
# For details, see the LICENSE file.
# Commercial use beyond a 30-day trial requires a separate license.
#
# Source Code: https://github.com/CoReason-AI/coreason_prism

from enum import Enum
from typing import Any, Dict, List, Optional
from uuid import UUID

from pydantic import BaseModel, Field, model_validator


class FigureType(str, Enum):
    """Type of extracted figure."""

    LINE_CHART = "LINE_CHART"
    BAR_CHART = "BAR_CHART"
    HISTOLOGY = "HISTOLOGY"
    UNKNOWN = "UNKNOWN"
    SKIPPED = "SKIPPED"


class ExtractedFigure(BaseModel):
    """Representation of an extracted figure/chart or bio-image.

    Attributes:
        id (UUID): Unique identifier.
        source_document_id (str): ID of the source document containing the figure.
        image_path (str): Path to the image file.
        figure_type (FigureType): Detected type of figure.
        data_series (Optional[Dict[str, List[Optional[float]]]]): Extracted data points
            (e.g. {'x': [...], 'y': [...]}).
        embedding (Optional[List[float]]): Multi-modal embedding (BioCLIP).
        metadata (Optional[Dict[str, Any]]): Additional metadata like median survival,
            cell counts, etc.
    """

    id: UUID
    source_document_id: str
    image_path: str
    figure_type: FigureType

    # The "De-plotted" Data
    # Allow Optional[float] to handle NaN/Missing values which serialize to null
    data_series: Optional[Dict[str, List[Optional[float]]]] = None  # { "x": [1,2,None], "y": [10, 20, 30] }

    # BioCLIP Embedding
    embedding: Optional[List[float]] = None

    # Biological Metadata (e.g., cell counts, ROI areas)
    metadata: Optional[Dict[str, Any]] = Field(default=None, description="Metadata from biological analysis")

    @model_validator(mode="after")
    def check_data_series_lengths(self) -> "ExtractedFigure":
        if self.data_series:
            lengths = [len(series) for series in self.data_series.values()]
            if len(set(lengths)) > 1:
                raise ValueError("All data series must have the same length")
        return self
