import numpy as np

from presupuestov1.schema import OrmBaseModel
from pydantic import field_serializer


class Treemap(OrmBaseModel):
    gestion: int
    entidad: int
    entidad_area: str | None
    nivel: str | None
    devengado: float | None

    @field_serializer('devengado')
    def do_round(self, v):
        if v is not None:
            return np.round(v, 2)


class TreemapObjeto(Treemap):
    objeto: int
    parent: int | None
    desc_objeto: str | None
    desc_entidad: str | None

class TreemapObjetoSingle(Treemap):
    objeto: int
    parent: int | None
    desc_entidad: str | None


class TreemapFinFun(Treemap):
    finfun: str
    parent: str | None
    desc_finfun: str | None
    desc_entidad: str | None


class TreemapFinFunSingle(Treemap):
    finfun: str
    parent: str | None
    desc_entidad: str | None


class TreemapActeco(Treemap):
    acteco: str
    parent: str | None
    desc_acteco: str | None
    desc_entidad: str | None


class TreemapActecoSingle(Treemap):
    acteco: str
    parent: str | None
    desc_entidad: str | None