from __future__ import annotations
from typing import List
import numpy as np
import logging
from silx.gui.plot3d.items.core import ItemChangedType
from silx.gui.plot3d.items import DataItem3D, ColormapMixIn
from silx.gui.plot3d.scene import primitives
from silx.io.url import DataUrl
from ewoks3dxrd.nexus.grains import read_grains
from ewoks3dxrd.nexus.utils import group_exists
from ImageD11.grain import read_grain_file
from ImageD11.grain import grain as Grain

_logger = logging.getLogger(__name__)


class Spheres(DataItem3D, ColormapMixIn):

    def __init__(self, positions=None, radii=None, values=None, parent=None):
        DataItem3D.__init__(self, parent=parent)
        ColormapMixIn.__init__(self)
        self._colorVal = values
        self._radiiNorm = 0.0001
        if positions is None:
            positions = np.zeros((0, 3), dtype=np.float32)
        if radii is None:
            radii = np.zeros((0, 1), dtype=np.float32)
        self._sphere = primitives.Spheres(positions, self._radiiNorm * radii)
        self._getScenePrimitive().children.append(self._sphere)
        self._colormapChanged()

    def _colormapChanged(self):
        super()._colormapChanged()
        self.setValues(self._colorVal, copy=False)

    def setRadii(self, radii: np.array):
        position = self._sphere.getAttribute(name="position", copy=False)
        if len(position) != len(radii):
            raise ValueError("Shape mismatch")
        self._sphere.setAttribute(name="radius", array=self._radiiNorm * radii)

    def setPosition(self, newPosition: np.array):
        oldPosition = self._sphere.getAttribute(name="position", copy=False)
        if newPosition.size != oldPosition.size:
            raise ValueError("Shape mismatch")
        self._sphere.setAttribute(name="position", array=newPosition)

    def setValues(self, values: np.array, copy: bool = True):
        self._colorVal = values
        if self._colorVal is None:
            self._sphere.setAttribute(name="color", array=(1.0, 1.0, 1.0, 1.0))
            return

        colors = self.getColormap().applyToData(self._colorVal) / 255
        self._sphere.setAttribute(name="color", array=colors, copy=True)

    def setRadiiNorm(self, radiiNorm: float):
        radii = self._sphere.getAttribute(name="radius") / self._radiiNorm
        self._radiiNorm = radiiNorm
        self.setRadii(radii=radii)

    def setData(self, positions: np.array, radii: np.array, values: np.array):
        self._sphere.setAttribute(name="position", array=positions)
        self._sphere.setAttribute(name="radius", array=self._radiiNorm * radii)
        self._colorVal = values
        colors = self.getColormap().applyToData(self._colorVal) / 255
        self._sphere.setAttribute(name="color", array=colors)
        self._updated(ItemChangedType.DATA)


def build_grain_spheres(grain_url: str, sphere_scale: float = 0.3) -> Spheres:
    data_url = DataUrl(grain_url)
    group_exists(filename=data_url.file_path(), data_group_path=data_url.data_path())
    grains = read_grains(
        grain_file_h5=data_url.file_path(),
        entry_name=data_url.data_path(),
        process_group_name="",
    )

    return build_spheres_from_grains_list(grains=grains, sphere_scale=sphere_scale)


def _parse_mean_intensity(intensity_info: str, default: float = 0.0) -> float:
    info_str = intensity_info.strip()
    if "no peaks" in info_str:
        return default

    try:
        mean_str = info_str.split("mean = ")[1].split(" , ")[0].replace("'", "")
        return float(mean_str)
    except (IndexError, ValueError) as e:
        print(
            f"Warning: Failed to parse intensity string '{intensity_info}'. Error: {e}"
        )
        return 0.0


def build_grains_from_ascii_file(
    grain_file: str, sphere_scale: float = 0.3
) -> Spheres | None:

    try:
        grains = read_grain_file(grain_file)
    except (IndexError, ValueError) as e:
        _logger.warning(
            f"Skipping reading UBI grain file. Error encountered while parsing file "
            f"'{grain_file}': {type(e).__name__} - {e}"
        )
        return None
    if not grains:
        return None

    return build_spheres_from_grains_list(grains=grains, sphere_scale=sphere_scale)


def build_spheres_from_grains_list(
    grains: List[Grain], sphere_scale: float = 0.01
) -> Spheres:
    position = np.array([grain.translation for grain in grains], dtype=np.float32)
    if hasattr(grains[0], "intensity_info"):
        mean_intensity = [
            _parse_mean_intensity(getattr(grain, "intensity_info", None))
            for grain in grains
        ]
    else:
        mean_intensity = [float(grain.mean_intensity) for grain in grains]

    radii = [m_int ** (1 / 3) for m_int in mean_intensity if m_int is not None]

    sphere = Spheres(
        positions=position,
        radii=np.array(radii),
        values=np.array([float(grain.npks) for grain in grains]),
    )
    sphere.setRadiiNorm(radiiNorm=sphere_scale)
    return sphere
