# generated by datamodel-codegen:
#   filename:  vfg_schema_0_2_0.json
#   timestamp: 2025-01-26T21:42:33+00:00

from __future__ import annotations

from enum import Enum
from typing import Dict, List, Optional, Union

from pydantic import BaseModel, Field, RootModel

# class Tensor(BaseModel):
#     __root__: Union[float, List[Tensor]]


class Tensor(RootModel[Union[float, List["Tensor"]]]):
    pass


class Distribution(str, Enum):
    Categorical = "categorical"
    CategoricalConditional = "categorical_conditional"


class Role(str, Enum):
    Likelihood = "likelihood"
    Transition = "transition"
    Preference = "preference"


class Factor(BaseModel):
    variables: List[str]
    distribution: Distribution
    values: Tensor
    role: Optional[Role] = None


class VFG(BaseModel):
    version: str = Field("0.2.0", Literal=True)
    factors: List[Factor]
    variables: Dict[str, List[str]]


Tensor.model_rebuild()


def migrate(data: dict) -> dict:
    return {
        "version": "0.4.0",
        "metadata": data.get("metadata", None),
        "variables": {
            name: {"elements": elements}
            for name, elements in data.get("variables", {}).items()
        },
        "factors": data.get("factors", None),
        "visualization_metadata": data.get("visualization_metadata", None),
    }
