from __future__ import annotations
from typing import ClassVar, TYPE_CHECKING
import numpy as np
from dimetra.geo.length_seq import LengthSeq
from h5kit import codec
from h5kit.codec import H5CodecLike, H5CodecMixin
from h5kit.io import H5Reader, H5Writer
from dimetra.h5.codec.length_seq_codec_layout import LENGTH_SEQ_CODEC_LAYOUT
from dimetra.geo.length_array import LengthArray
from dimetra.geo.length_unit import length_unit_from_token

if TYPE_CHECKING:
    from h5kit.codec.h5_codec_layout import H5CodecLayout

@codec
class LengthSeqCodec(H5CodecMixin, H5CodecLike[LengthSeq]):
    PY_TYPE: ClassVar[type[LengthSeq]] = LengthSeq
    SCHEMA: ClassVar[str] = "allytools.units.LengthSeq"
    VER: ClassVar[int] = 1
    LAYOUT: ClassVar[H5CodecLayout] = LENGTH_SEQ_CODEC_LAYOUT.h5_layout
    DEFAULT_UNIT: ClassVar[str] = "mm"

    @classmethod
    def write(cls, w: H5Writer, obj: LengthSeq) -> None:
        cls.clear_before_write(w)
        values_mm = np.asarray(obj.values_mm, dtype=np.float64)
        unit_token = obj.unit.name.lower()
        w.array(LENGTH_SEQ_CODEC_LAYOUT.values_mm, values_mm)
        w.attr(LENGTH_SEQ_CODEC_LAYOUT.unit, unit_token)

    @classmethod
    def read(cls, r: H5Reader) -> LengthSeq:
        values_mm = np.asarray(r.require_ndarray(LENGTH_SEQ_CODEC_LAYOUT.values_mm),dtype=np.float64)
        unit_token = str(r.attr(LENGTH_SEQ_CODEC_LAYOUT.unit, cls.DEFAULT_UNIT))
        unit = length_unit_from_token(unit_token)
        data = LengthArray(values_mm=values_mm, unit=unit)
        return LengthSeq(data)
