"""Scene composition layer for ModernGL renders.

Scenes combine one or more low-level render objects into a user-facing
configuration. A scene owns the system being shown, scene-level state,
property groups, camera-slot policy, and the orchestration needed to upload and
render all child renders consistently.
"""

from __future__ import annotations

import moderngl as mgl
import numpy as np

from ..core import System
from .buffer import Field3DBuffer, Field1DBuffer, IntBuffer
from .sphere import SphereDotRender
from .vector import VectorRender
from .charge import ChargeRender
from .flow import FlowRender, AstraRender
from .isolines import IsolinesRender
from .camera import Camera
from .render import _release_nested_resources
from ..core import (
    NumericProperty,
    BooleanProperty,
    ListProperty,
    PropertiesCollection,
    UpdateMask,
)
from .level import LevelRender, ScalarLevelRender
from .axes import AxesRender
from .target import RenderTarget


class Scene:
    """Compose one or more render units into a backend-agnostic scene.

    Scenes are the main reusable rendering abstraction in TopoVec. They expose
    scene-level properties, choose default cameras for logical camera slots, and
    hide the details of how many low-level render objects participate in the
    final frame.
    """

    def __init__(
        self,
        mglctx: mgl.Context,
        system: System,
        camera: Camera = None,
        state: dict[str, object] | None = None,
    ):
        """Create a scene bound to one OpenGL context and one lattice system."""
        self._system = system
        self._ctx = mglctx
        self._initial_state = {} if state is None else dict(state)
        self._provided_camera = camera
        self._camera = (
            self.__class__.make_default_camera_for(system, state=self._initial_state)
            if camera is None
            else camera
        )
        self._cameras = {
            str(self.__class__.camera_slot_key_for_state(self._initial_state)): self._camera
        }
        assert isinstance(self._system, System)
        assert isinstance(self._ctx, mgl.Context)
        assert isinstance(self._camera, Camera)

    def __getitem__(self, key):
        """Return one property group as a mutable :class:`PropertiesCollection`."""
        props_by_render = self.list_properties()
        if key not in props_by_render:
            raise KeyError(
                f"Render `{key}` is unknown. To see list of renders call `yourscene.print_properties()`."
            )
        return PropertiesCollection(props_by_render[key])

    def upload(self, state: np.ndarray = None):
        """Upload the field/state consumed by the scene."""
        raise NotImplementedError

    def upload_mask(self, mask: np.ndarray = None):
        """Specify grid points to show.

        Renders showing individual vectors/directors at grid points can remove some of the points
        according to the mask.
        """
        pass

    def render(self, aspectratio=None):
        """Compatibility wrapper rendering into the current framebuffer."""
        target = RenderTarget.from_current(
            self._ctx,
            aspectratio=aspectratio,
        )
        return self.render_to(target)

    def render_to(self, target: RenderTarget):
        """Render the scene into one explicit framebuffer target."""
        legacy_render = type(self).render
        if legacy_render is Scene.render:
            raise NotImplementedError
        return legacy_render(self, aspectratio=target.aspectratio)

    def clear_background(self, target: RenderTarget | None = None, color=None):
        """Clear the current framebuffer using an opaque background by default."""
        resolved = self._background if color is None else color
        render_target = (
            RenderTarget.from_current(self._ctx)
            if target is None
            else target
        )
        render_target.clear_color(resolved)

    def clear_current_framebuffer(self, *rgba: float) -> None:
        """Clear the currently bound framebuffer explicitly."""
        RenderTarget.from_current(self._ctx).clear_color(rgba)

    @property
    def camera(self) -> Camera:
        """Common camera for all renders."""
        return self._camera

    @property
    def system(self) -> System:
        """System rendered by the scene."""
        return self._system

    @property
    def camera_slot_key(self) -> str:
        """Logical camera slot shared across compatible scenes."""
        return str(self.__class__.camera_slot_key_for_state(self.camera_state()))

    @property
    def camera_profile(self) -> str:
        """Backward-compatible alias for `camera_slot_key`."""
        return self.camera_slot_key

    def set_camera(self, camera, slot: str | None = None):
        """Replace the camera object for the active or explicit slot."""
        slot_name = str(self.camera_slot_key if slot is None else slot)
        self._cameras[slot_name] = camera
        self._camera = camera

    @classmethod
    def camera_slot_key_for_state(cls, state: dict[str, object] | None = None) -> str:
        """Return the logical camera slot for a scene state snapshot."""
        return "general"

    @classmethod
    def make_default_camera_for(
        cls,
        system: System,
        state: dict[str, object] | None = None,
    ) -> Camera:
        """Create the default camera for a system/state pair."""
        camera = Camera()
        camera.adjust_fov(system)
        camera.adjust_scale(system)
        return camera

    def make_default_camera(self) -> Camera:
        """Create the default camera for the current live scene state."""
        return self.__class__.make_default_camera_for(
            self._system,
            state=self.camera_state(),
        )

    def camera_state(self) -> dict[str, object]:
        """Return the scene state relevant for camera-slot selection."""
        return dict(self._initial_state)

    def activate_camera_slot(self, slot: str | None = None) -> Camera:
        """Activate or lazily create one logical camera slot."""
        slot_name = str(self.camera_slot_key if slot is None else slot)
        if slot_name not in self._cameras:
            self._cameras[slot_name] = self.make_default_camera()
        self._camera = self._cameras[slot_name]
        return self._camera

    def _apply_initial_state(self) -> None:
        """Apply constructor-provided scene state after child renders exist."""
        for field, value in self._initial_state.items():
            if not hasattr(self, field):
                raise AttributeError(
                    f"Scene {self.__class__.__name__} has no state field {field!r}."
                )
            setattr(self, field, value)

    def _finalize_scene_init(self) -> None:
        """Finalize camera/state initialization in subclasses.

        Leaf scenes call this after all render children are constructed.
        """
        self._apply_initial_state()
        slot_name = self.camera_slot_key
        if self._provided_camera is not None:
            self._cameras = {slot_name: self._provided_camera}
            self._camera = self._provided_camera
            return
        self._camera = self.make_default_camera()
        self._cameras[slot_name] = self._camera

    @classmethod
    def supports_system(cls, system: System) -> bool:
        """Report whether the scene contract is compatible with ``system``."""
        return True

    def list_properties(self):
        """Return property groups exposed by the scene and its child renders."""
        return {}

    def list_actions(self):
        """Get all available actions."""
        return {}

    def release(self) -> None:
        """Release child GL resources and retained large arrays."""
        for name, value in list(getattr(self, "__dict__", {}).items()):
            if name in {
                "_ctx",
                "_system",
                "_camera",
                "_cameras",
                "_provided_camera",
                "_initial_state",
            }:
                continue
            if isinstance(value, np.ndarray):
                self.__dict__[name] = None
                continue
            _release_nested_resources(value)

    def print_properties(self):
        props_by_render = self.list_properties()
        for render, props in props_by_render.items():
            print(f"'{render:}'")
            for n, prop in enumerate(props):
                print(f"  {n}: {prop}")

    @property
    def detalization(self):
        """Detalization level.

        Higher levels of detalization give smoother surfaces, but render much slower.
        """
        return 1.0

    @detalization.setter
    def detalization(self, value):
        pass


##################################################################################################


class EmptyScene(Scene):
    """Base scene with background clearing and optional coordinate axes."""

    def __init__(self, background=(1, 1, 1), *vargs, **kwargs):
        super().__init__(*vargs, **kwargs)
        self._background = tuple(float(value) for value in background[:3])
        self._white_background = bool(
            np.allclose(self._background, (1.0, 1.0, 1.0))
        )
        self.preferable_size = np.min(np.linalg.norm(self._system.primitives, axis=-1))

        self._axes = AxesRender(mglctx=self._ctx)
        self.show_axes = True

    @property
    def white_background(self) -> bool:
        return bool(self._white_background)

    @white_background.setter
    def white_background(self, value: bool) -> None:
        self._white_background = bool(value)
        self._background = (
            (1.0, 1.0, 1.0) if self._white_background else (0.0, 0.0, 0.0)
        )

    def upload(self, state=None, **kwargs):
        pass

    def render_to(self, target: RenderTarget):
        target.bind()
        self.clear_background(target)

    @property
    def axes_render(self):
        return self._axes

    def _style_properties(self) -> list[BooleanProperty]:
        return [
            BooleanProperty(
                self,
                "white_background",
                title="White background",
                param_key="style.white_background",
            ),
        ]

    def _axes_properties(self) -> list[BooleanProperty]:
        return [
            BooleanProperty(
                self,
                "show_axes",
                title="Show axes",
                param_key="scene.show_axes",
            ),
        ] + self._axes.list_properties()

    def render_axes(self, target: RenderTarget):
        if not self.show_axes:
            return
        self._axes.render(
            modelview=self.camera.get_view_matrix(),
            projection=self.camera.get_projection_matrix(aspectratio=target.aspectratio),
            target=target,
        )

    def list_properties(self):
        return {
            "Style": self._style_properties(),
            "Axes": self._axes_properties(),
            **super().list_properties(),
        }
    

#################################################################################


class OrthoScene(EmptyScene):
    """Common ancestor for 3D scenes preferring orthographic projection."""

    @classmethod
    def supports_system(cls, system: System) -> bool:
        return int(getattr(system, "dim", 0)) == 3

    @classmethod
    def make_default_camera_for(
        cls,
        system: System,
        state: dict[str, object] | None = None,
    ) -> Camera:
        camera = super().make_default_camera_for(system, state=state)
        camera.perspective_is_on = False
        return camera

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    def compute_camera_basis(self, axis):
        """Compute natural basis for camera looking along `axis`.

        Similar functionality is given by `Camera.set_predefined_direction`.
        """
        return self.__class__.compute_camera_basis_for(self._system, axis)

    @classmethod
    def compute_camera_basis_for(cls, system, axis):
        assert system.dim == 3
        size = system.size
        primitives = system.primitives
        if primitives.shape[0] == 3:
            basis = primitives
        elif primitives.shape[0] == 2:
            basis = np.empty((3, 3), dtype=np.float32)
            basis[:2] = primitives
            basis[2] = np.cross(basis[0], basis[1])
        else:
            assert False

        bxis = (axis - 1) % 3
        cxis = (axis - 2) % 3
        if size[bxis] < size[cxis]:
            depth = basis[axis]
            up = -basis[bxis]
        else:
            depth = basis[axis]
            up = basis[cxis]
        return (depth, up)


##########################################################################################################


class SphereDotScene(EmptyScene):
    """Render all directors on Block sphere."""

    @classmethod
    def camera_slot_key_for_state(cls, state: dict[str, object] | None = None) -> str:
        return "bloch"

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._buffer = Field3DBuffer(self._ctx, self._system)
        self._render = SphereDotRender(
            mglctx=self._ctx, system=self._system, stateBuffer=self._buffer
        )
        self._finalize_scene_init()

    def upload(self, state: np.ndarray = None, **kwargs):
        self._buffer.upload(state)

    def render_to(self, target: RenderTarget):
        target.bind()
        self.clear_background(target)
        self._ctx.enable(mgl.DEPTH_TEST)
        self._render.render(
            modelview=self.camera.get_view_matrix(),
            projection=self.camera.get_projection_matrix(aspectratio=target.aspectratio),
        )

    def list_properties(self):
        return {
            "Ball": self._render.list_properties(),
            **super().list_properties(),
        }


##########################################################################################################


class VectorScene(EmptyScene):
    """Render director field as vectors/cylinders at every grid point."""

    @classmethod
    def supports_energy_mask_system(cls, system: System) -> bool:
        return callable(getattr(system, "field3D", None)) and callable(
            getattr(system, "field1D", None)
        )

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._state = None
        self._m_state = self._system.field3D() if self.supports_energy_mask else None
        self._buffer = Field3DBuffer(mglctx=self._ctx, system=self._system)
        self._coordinates_buffer = Field3DBuffer(mglctx=self._ctx, system=self._system)
        self._mask_buffer = Field3DBuffer(mglctx=self._ctx, system=self._system)
        self._scalar_mask_buffer = Field1DBuffer(mglctx=self._ctx, system=self._system)

        self._render = VectorRender(
            mglctx=self._ctx,
            system=self._system,
            stateBuffer=self._buffer,
            coordinatesBuffer=self._coordinates_buffer,
            maskBuffer=self._mask_buffer,
            scalarMaskBuffer=self._scalar_mask_buffer,
        )
        self._render.base_size = self.preferable_size
        self._finalize_scene_init()

    @property
    def vector_render(self):
        return self._render

    @property
    def supports_energy_mask(self) -> bool:
        return self.__class__.supports_energy_mask_system(self._system)

    def get_use_energy_mask(self):
        return self._render.use_scalar_mask

    def set_use_energy_mask(self, value):
        self._render.use_scalar_mask = bool(value) and self.supports_energy_mask
        self.update_scalar_mask()
        return UpdateMask.REDRAW

    use_energy_mask = property(get_use_energy_mask, set_use_energy_mask)

    def update_scalar_mask(self):
        # print("update_scalar_mask")
        if not self.supports_energy_mask:
            self._render.use_scalar_mask = False
            return
        if not self._render.use_scalar_mask:
            return
        if self._state is None:
            return
        self._render.use_mask = True

        m_scalar1 = self._system.field1D()
        m_scalar2 = self._system.field1D()
        self._m_state.upload(self._state).energy_contributions(
            dm=m_scalar1, heisenberg=m_scalar2
        )
        scalar = -m_scalar1.download() - m_scalar2.download()
        mn, mx = np.min(scalar), np.max(scalar)
        scalar = (scalar - mn) / (mx - mn)
        scalar = np.power(scalar, 32)
        # scalar = (np.exp(scalar)-1)/np.exp(1)

        # scalar = -self._m_state.upload(self._state).energy().download()
        # scalar = np.abs(self._m_state.upload(self._state).top_charge().download())
        self._scalar_mask_buffer.upload(scalar)
        self._render.min_scalar = 0  # np.min(scalar)
        self._render.max_scalar = 1  # np.max(scalar)
        # print(self._render.min_scalar, self._render.max_scalar, self._render.use_scalar_mask)

    def upload(self, state: np.ndarray = None):
        self._state = state
        if not state is None:
            self._buffer.upload(state)
            self.update_scalar_mask()

    def upload_mask(self, mask: np.ndarray = None):
        self._render.use_scalar_mask = False
        if mask is None:
            self._render.use_mask = False
        else:
            self._mask_buffer.upload(mask)
            self._render.use_mask = True

    def render_to(self, target: RenderTarget):
        target.bind()
        self.clear_background(target)
        self._ctx.enable_only(mgl.DEPTH_TEST)
        self._render.render(
            modelview=self.camera.get_view_matrix(),
            projection=self.camera.get_projection_matrix(aspectratio=target.aspectratio),
            camera_position=self.camera.camera_position(),
        )

        self.render_axes(target)

    def list_properties(self):
        groups = {
            "Cones": self._render.list_properties(),
            **super().list_properties(),
        }
        if self.supports_energy_mask:
            groups["Mask"] = [
                BooleanProperty(
                    self,
                    "use_energy_mask",
                    title="Energy mask",
                    reset=True,
                ),
            ]
        return groups

    @property
    def detalization(self):
        return self._render.detalization

    @detalization.setter
    def detalization(self, value):
        self._render.detalization = value


##########################################################################################################


class LayerScene(OrthoScene):
    """Show sections of director field by a coordinate axis.

    Very similar to `VectorScene` but show only gird points sharing one coordinate defined by `axis` property.
    """

    DEFAULT_AXIS = 2

    @classmethod
    def camera_slot_key_for_state(cls, state: dict[str, object] | None = None) -> str:
        axis = int((state or {}).get("axis", cls.DEFAULT_AXIS))
        axis_name = ("x", "y", "z")[axis]
        return f"layer-{axis_name}"

    @classmethod
    def make_default_camera_for(
        cls,
        system: System,
        state: dict[str, object] | None = None,
    ) -> Camera:
        axis = int((state or {}).get("axis", cls.DEFAULT_AXIS))
        camera = Camera(perspective=False)
        camera.reset_center_of_FOV(system)
        camera.set_basis(*cls.compute_camera_basis_for(system, axis))
        camera.adjust_scale(system)
        return camera

    def __init__(self, finalize: bool = True, **kwargs):
        super().__init__(**kwargs)
        self._buffer = Field3DBuffer(mglctx=self._ctx, system=self._system)
        self._coordinates_buffer = Field3DBuffer(mglctx=self._ctx, system=self._system)
        self._mask_buffer = self._coordinates_buffer  # Not used
        self._scalar_mask_buffer = self._coordinates_buffer

        self._render = VectorRender(
            mglctx=self._ctx,
            system=self._system,
            stateBuffer=self._buffer,
            coordinatesBuffer=self._coordinates_buffer,
            maskBuffer=self._mask_buffer,
            scalarMaskBuffer=self._scalar_mask_buffer,
        )
        self._render.min_size = 0
        self._render.base_size = self.preferable_size
        self._render.min_threshold = 0
        self._render.power_size = 1

        self.axis = self.DEFAULT_AXIS
        self.show_axes = True
        if finalize:
            self._finalize_scene_init()

    def camera_state(self) -> dict[str, object]:
        return {
            "axis": int(
                getattr(self, "_axis", self._initial_state.get("axis", self.DEFAULT_AXIS))
            )
        }

    @property
    def vector_render(self):
        return self._render

    def get_axis(self):
        return self._axis

    def set_axis(self, value):
        self._axis = int(value)
        self.layers_count = self._render.size[self._axis]
        self.layer = self.layers_count / 2
        return (
            UpdateMask.REDRAW
            | UpdateMask.REBUILD_PROPERTIES
            | UpdateMask.CAMERA_SLOT_CHANGED
        )

    def get_layer(self):
        return self._layer

    def set_layer(self, value):
        self._layer = int(value)
        mn = 0 * np.array(self._render.size)
        mx = np.array(self._render.size)
        mn[self._axis] = self._layer
        mx[self._axis] = self._layer
        self._render.clip_min = mn
        self._render.clip_max = mx
        return UpdateMask.REDRAW

    axis = property(get_axis, set_axis)
    layer = property(get_layer, set_layer)

    def upload(self, state: np.ndarray = None):
        if not state is None:
            self._buffer.upload(state)

    def upload_mask(self, mask: np.ndarray = None):
        pass

    def render_to(self, target: RenderTarget):
        target.bind()
        self.clear_background(target)
        self._ctx.enable_only(mgl.DEPTH_TEST)
        self._render.render(
            modelview=self.camera.get_view_matrix(),
            projection=self.camera.get_projection_matrix(aspectratio=target.aspectratio),
            camera_position=self.camera.camera_position(),
        )

        self.render_axes(target)

    def list_properties(self):
        return {
            "Filter": [
                NumericProperty(
                    self,
                    "axis",
                    title="Axis",
                    min=0,
                    max=2,
                    count=3,
                    reset=True,
                    param_key="layer.axis",
                ),
                NumericProperty(
                    self,
                    "layer",
                    title="Layer",
                    min=0,
                    max=self.layers_count - 1,
                    count=self.layers_count,
                    param_key="layer.index",
                ),
            ],
            # "Cones": self._render.list_properties(),
            "Cones": [
                ListProperty(
                    self._render,
                    "mesh_shape",
                    title="Shape",
                    values=sorted(self._render.MESH_NAMES.keys()),
                    param_key="layer.mesh_shape",
                ),
                # BooleanProperty(self._render, "director_mode", title="Director mode"),
                BooleanProperty(self._render, "negate", title="Negate spin"),
                NumericProperty(
                    self._render,
                    "width",
                    title="Width",
                    min=0.01,
                    max=2.0,
                    count=100,
                    param_key="layer.vector_width",
                ),
                NumericProperty(
                    self._render,
                    "color_shift",
                    title="Hue",
                    min=0.0,
                    max=1.0,
                    count=360,
                ),
                NumericProperty(
                    self._render,
                    "saturation_mag",
                    title="Saturation",
                    min=1.0,
                    max=10.0,
                    count=500,
                ),
                # NumericProperty(self, "min_size", title="Min. size", min=0.0, max=2.0, count=30),
                NumericProperty(
                    self._render,
                    "max_size",
                    title="Max. size",
                    min=0.0,
                    max=2.0,
                    count=30,
                    param_key="layer.vector_max_size",
                ),
                # NumericProperty(self._render, "power_size", title="Steepness", min=0.0, max=2.0, count=30),
                # NumericProperty(self._render, "min_threshold", title="Min. threshold", min=0.0, max=1.0, count=500),
                # NumericProperty(self._render, "detalization", title="Detalization", min=1., max=10., count=10),
            ],
            **super().list_properties(),
        }

    @property
    def detalization(self):
        return self._render.detalization

    @detalization.setter
    def detalization(self, value):
        self._render.detalization = value


##########################################################################################################


class LayerAndSurfaceScene(LayerScene):
    """Mix of LayerScene and SurfaceScene."""

    @classmethod
    def make_default_camera_for(
        cls,
        system: System,
        state: dict[str, object] | None = None,
    ) -> Camera:
        camera = super().make_default_camera_for(system, state=state)
        camera.perspective_is_on = True
        return camera

    def __init__(self, **kwargs):
        super().__init__(finalize=False, **kwargs)

        idx = self._system.get_inner_simplices()
        self._idx = IntBuffer(mglctx=self._ctx, count=4 * idx.shape[0])
        self._idx.upload(idx.flatten())

        self._level_render = FlowRender(
            mglctx=self._ctx,
            system=self._system,
            stateBuffer=self._buffer,
            indexBuffer=self._idx,
            coordinatesBuffer=self._coordinates_buffer,
        )

        self._finalize_scene_init()

    @property
    def vector_render(self):
        return self._render

    @property
    def isosurface_render(self):
        return self._level_render

    def render_to(self, target: RenderTarget):
        target.bind()
        self.clear_background(target)
        self._ctx.enable_only(mgl.DEPTH_TEST)

        self._render.color_shift = self._level_render.color_shift
        self._render.lighting = self._level_render.lighting
        self._render.negate = self._level_render.negate

        self._render.render(
            modelview=self.camera.get_view_matrix(),
            projection=self.camera.get_projection_matrix(aspectratio=target.aspectratio),
            camera_position=self.camera.camera_position(),
        )

        self._level_render.render(
            modelview=self.camera.get_view_matrix(),
            projection=self.camera.get_projection_matrix(aspectratio=target.aspectratio),
            camera_position=self.camera.camera_position(),
        )

        self.render_axes(target)

    def list_properties(self):
        base_groups = super().list_properties()
        style_properties = [
            *self._style_properties(),
            BooleanProperty(
                self._level_render,
                "lighting",
                title="Lighting",
                param_key="style.lighting",
            ),
            BooleanProperty(
                self._level_render,
                "negate",
                title="Negate",
                param_key="style.negate",
            ),
            NumericProperty(
                self._level_render,
                "color_shift",
                title="Hue",
                min=0.0,
                max=1.0,
                count=360,
                param_key="style.color_shift",
            ),
        ]
        axes_group = base_groups.pop("Axes")
        base_groups.pop("Style", None)
        return {
            "Layer": [
                NumericProperty(
                    self,
                    "axis",
                    title="Axis",
                    min=0,
                    max=2,
                    count=3,
                    reset=True,
                    param_key="layer.axis",
                ),
                NumericProperty(
                    self,
                    "layer",
                    title="Layer",
                    min=0,
                    max=self.layers_count - 1,
                    count=self.layers_count,
                    param_key="layer.index",
                ),
                ListProperty(
                    self._render,
                    "mesh_shape",
                    title="Shape",
                    values=sorted(self._render.MESH_NAMES.keys()),
                    param_key="layer.mesh_shape",
                ),
                NumericProperty(
                    self._render,
                    "width",
                    title="Width",
                    min=0.01,
                    max=2.0,
                    count=100,
                    param_key="layer.vector_width",
                ),
                NumericProperty(
                    self._render,
                    "max_size",
                    title="Max. size",
                    min=0.0,
                    max=2.0,
                    count=30,
                    param_key="layer.vector_max_size",
                ),
            ],
            "Isosurface": [
                BooleanProperty(
                    self._level_render,
                    "neglect_z",
                    title="Ignore z-projection",
                    param_key="surface.neglect_z",
                ),
                ListProperty(
                    self._level_render,
                    "angle_name",
                    title="Constant",
                    values=self._level_render.ANGLENAMES,
                    param_key="surface.angle_name",
                ),
                NumericProperty(
                    self._level_render,
                    "level",
                    title="Level",
                    min=-1.0,
                    max=1.0,
                    count=1000,
                    param_key="surface.level",
                ),
            ],
            "Style": style_properties,
            **base_groups,
            "Axes": axes_group,
        }


##########################################################################################################


class LayerAndIsolinesScene(LayerScene):
    """Mix of LayerScene and IsolinesScene."""

    @classmethod
    def make_default_camera_for(
        cls,
        system: System,
        state: dict[str, object] | None = None,
    ) -> Camera:
        camera = super().make_default_camera_for(system, state=state)
        camera.perspective_is_on = True
        return camera

    def __init__(self, **kwargs):
        super().__init__(finalize=False, **kwargs)

        idx = self._system.get_inner_simplices()
        self._idx = IntBuffer(mglctx=self._ctx, count=4 * idx.shape[0])
        self._idx.upload(idx.flatten())

        self._level_render = IsolinesRender(
            mglctx=self._ctx,
            system=self._system,
            stateBuffer=self._buffer,
            indexBuffer=self._idx,
            coordinatesBuffer=self._coordinates_buffer,
        )

        self._finalize_scene_init()

    @property
    def vector_render(self):
        return self._render

    @property
    def isolines_render(self):
        return self._level_render

    def render_to(self, target: RenderTarget):
        target.bind()
        self.clear_background(target)
        self._ctx.enable_only(mgl.DEPTH_TEST)

        self._render.color_shift = self._level_render.color_shift
        self._render.lighting = self._level_render.lighting
        self._render.negate = self._level_render.negate

        self._render.render(
            modelview=self.camera.get_view_matrix(),
            projection=self.camera.get_projection_matrix(aspectratio=target.aspectratio),
            camera_position=self.camera.camera_position(),
        )

        self._ctx.line_width = 3
        self._level_render.render(
            modelview=self.camera.get_view_matrix(),
            projection=self.camera.get_projection_matrix(aspectratio=target.aspectratio),
            camera_position=self.camera.camera_position(),
        )

        self.render_axes(target)

    def list_properties(self):
        base_groups = super().list_properties()
        style_properties = [
            *self._style_properties(),
            BooleanProperty(
                self._level_render,
                "lighting",
                title="Lighting",
                param_key="style.lighting",
            ),
            BooleanProperty(
                self._level_render,
                "negate",
                title="Negate",
                param_key="style.negate",
            ),
            NumericProperty(
                self._level_render,
                "color_shift",
                title="Hue",
                min=0.0,
                max=1.0,
                count=360,
                param_key="style.color_shift",
            ),
        ]
        axes_group = base_groups.pop("Axes")
        base_groups.pop("Style", None)
        return {
            "Layer": [
                NumericProperty(
                    self,
                    "axis",
                    title="Axis",
                    min=0,
                    max=2,
                    count=3,
                    reset=True,
                    param_key="layer.axis",
                ),
                NumericProperty(
                    self,
                    "layer",
                    title="Layer",
                    min=0,
                    max=self.layers_count - 1,
                    count=self.layers_count,
                    param_key="layer.index",
                ),
                ListProperty(
                    self._render,
                    "mesh_shape",
                    title="Shape",
                    values=sorted(self._render.MESH_NAMES.keys()),
                    param_key="layer.mesh_shape",
                ),
                NumericProperty(
                    self._render,
                    "width",
                    title="Width",
                    min=0.01,
                    max=2.0,
                    count=100,
                    param_key="layer.vector_width",
                ),
                NumericProperty(
                    self._render,
                    "max_size",
                    title="Max. size",
                    min=0.0,
                    max=2.0,
                    count=30,
                    param_key="layer.vector_max_size",
                ),
            ],
            "Isolines": [
                BooleanProperty(
                    self._level_render,
                    "neglect_z",
                    title="Ignore z-projection",
                    param_key="surface.neglect_z",
                ),
                NumericProperty(
                    self._level_render,
                    "phi_shift",
                    title="Polar angle offset",
                    min=-1.0,
                    max=1.0,
                    count=1000,
                    param_key="isolines.phi_shift",
                ),
                NumericProperty(
                    self._level_render,
                    "theta_shift",
                    title="Azimuthal angle offset",
                    min=-1.0,
                    max=1.0,
                    count=1000,
                    param_key="isolines.theta_shift",
                ),
                NumericProperty(
                    self._level_render,
                    "npetals",
                    title="Polar angle substeps",
                    min=1,
                    max=10,
                    count=10,
                    param_key="isolines.npetals",
                ),
                NumericProperty(
                    self._level_render,
                    "ntheta",
                    title="Azimuthal angle substeps",
                    min=1,
                    max=5,
                    count=5,
                    param_key="isolines.ntheta",
                ),
                NumericProperty(
                    self._level_render,
                    "width",
                    title="Width",
                    min=0.0,
                    max=0.02,
                    count=1000,
                    param_key="isolines.width",
                ),
            ],
            "Style": style_properties,
            **base_groups,
            "Axes": axes_group,
        }


##########################################################################################################


class CubeScene(OrthoScene):
    """Show director field as vectors similar to VectorScene with cubic cut."""

    @classmethod
    def make_default_camera_for(
        cls,
        system: System,
        state: dict[str, object] | None = None,
    ) -> Camera:
        state = {} if state is None else dict(state)
        reverse = np.array(
            [
                bool(state.get("reverse_x", False)),
                bool(state.get("reverse_y", False)),
                bool(state.get("reverse_z", False)),
            ]
        )
        camera = Camera(perspective=False)
        camera.reset_center_of_FOV(system)
        camera.set_basis(*cls.compute_camera_basis_for(system, reverse))
        camera.adjust_scale(system)
        return camera

    def __init__(self, *vargs, **kwargs):
        super().__init__(*vargs, **kwargs)
        self._buffer = Field3DBuffer(mglctx=self._ctx, system=self._system)
        self._coordinates_buffer = Field3DBuffer(mglctx=self._ctx, system=self._system)
        self._mask_buffer = self._coordinates_buffer  # Not used
        self._scalar_mask_buffer = self._coordinates_buffer

        self._render = VectorRender(
            mglctx=self._ctx,
            system=self._system,
            stateBuffer=self._buffer,
            coordinatesBuffer=self._coordinates_buffer,
            maskBuffer=self._mask_buffer,
            scalarMaskBuffer=self._scalar_mask_buffer,
        )
        self._render.min_size = 0
        self._render.base_size = self.preferable_size
        self._render.min_threshold = 0
        self._render.power_size = 1
        self._render.clip_invert = True

        self._section = np.array(self._render.size, dtype=np.int32) / 2
        self._reverse = np.array(
            [
                False,
            ]
            * 3
        )
        self.update_parameters()
        self._finalize_scene_init()

    def camera_state(self) -> dict[str, object]:
        return {
            "reverse_x": bool(getattr(self, "_reverse", [False, False, False])[0]),
            "reverse_y": bool(getattr(self, "_reverse", [False, False, False])[1]),
            "reverse_z": bool(getattr(self, "_reverse", [False, False, False])[2]),
        }

    def update_parameters(self):
        self._render.clip_min = [
            0 if self._reverse[n] else self._section[n] for n in range(3)
        ]
        self._render.clip_max = [
            self._section[n] + 1 if self._reverse[n] else self._render.size[n]
            for n in range(3)
        ]

    def get_section(self, n):
        return self._section[n]

    def set_section(self, n, value):
        self._section[n] = int(value)
        self.update_parameters()
        return UpdateMask.REDRAW

    def get_reverse(self, n):
        return self._reverse[n]

    def set_reverse(self, n, value):
        self._reverse[n] = bool(value)
        self.update_parameters()
        return UpdateMask.REDRAW

    section_x = property(
        lambda self: self.get_section(0), lambda self, v: self.set_section(0, v)
    )
    section_y = property(
        lambda self: self.get_section(1), lambda self, v: self.set_section(1, v)
    )
    section_z = property(
        lambda self: self.get_section(2), lambda self, v: self.set_section(2, v)
    )

    reverse_x = property(
        lambda self: self.get_reverse(0), lambda self, v: self.set_reverse(0, v)
    )
    reverse_y = property(
        lambda self: self.get_reverse(1), lambda self, v: self.set_reverse(1, v)
    )
    reverse_z = property(
        lambda self: self.get_reverse(2), lambda self, v: self.set_reverse(2, v)
    )

    def compute_camera_basis(self):
        return self.__class__.compute_camera_basis_for(self._system, self._reverse)

    @classmethod
    def compute_camera_basis_for(cls, system, reverse):
        d = np.array([1 if reverse[n] else -1 for n in range(3)])
        primitives = system.primitives

        if primitives.shape[0] == 3:
            basis = primitives
        elif primitives.shape[0] == 2:
            basis = np.empty((3, 3), dtype=np.float32)
            basis[:2] = primitives
            basis[2] = np.cross(basis[0], basis[1])
        else:
            assert False

        depth = np.sum(basis * d[:, None], axis=-1)
        d[1] *= -1
        up = np.sum(basis * d[:, None], axis=-1)
        return (depth, -up)

    def upload(self, state: np.ndarray = None):
        if not state is None:
            self._buffer.upload(state)

    def upload_mask(self, mask: np.ndarray = None):
        pass

    def render_to(self, target: RenderTarget):
        target.bind()
        self.clear_background(target)
        self._ctx.enable_only(mgl.DEPTH_TEST)
        self._render.render(
            modelview=self.camera.get_view_matrix(),
            projection=self.camera.get_projection_matrix(aspectratio=target.aspectratio),
            camera_position=self.camera.camera_position(),
        )

        self.render_axes(target)

    def list_properties(self):
        return {
            "Filter": [
                NumericProperty(
                    self,
                    "section_x",
                    title="X-section",
                    min=0,
                    max=self._render.size[0],
                    count=self._render.size[0] + 1,
                ),
                NumericProperty(
                    self,
                    "section_y",
                    title="Y-section",
                    min=0,
                    max=self._render.size[1],
                    count=self._render.size[1] + 1,
                ),
                NumericProperty(
                    self,
                    "section_z",
                    title="Z-section",
                    min=0,
                    max=self._render.size[2],
                    count=self._render.size[2] + 1,
                ),
                BooleanProperty(self, "reverse_x", title="Reverse X"),
                BooleanProperty(self, "reverse_y", title="Reverse Y"),
                BooleanProperty(self, "reverse_z", title="Reverse Z"),
            ],
            "Cones": [
                NumericProperty(
                    self._render, "max_size", title="Size", min=0.0, max=2.0, count=30
                ),
            ],
            **super().list_properties(),
        }

    @property
    def detalization(self):
        return self._render.detalization

    @detalization.setter
    def detalization(self, value):
        self._render.detalization = value


##########################################################################################################


class ChargeScene(OrthoScene):
    """Show distribution of topological charge.

    Charge computed on each coordinate plane in shown by RGB components of color.
    """

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._buffer = Field3DBuffer(mglctx=self._ctx, system=self._system)
        self._coordinates_buffer = Field3DBuffer(mglctx=self._ctx, system=self._system)

        idx = self._system.get_inner_triangles(
            triangulation=self._system.default_trivial_cell_triangulation(
                self._system.dim
            )
        )
        self._idx = IntBuffer(mglctx=self._ctx, count=3 * idx.shape[0])
        self._idx.upload(idx.flatten())

        self._render = ChargeRender(
            mglctx=self._ctx,
            system=self._system,
            stateBuffer=self._buffer,
            indexBuffer=self._idx,
            coordinatesBuffer=self._coordinates_buffer,
        )

        self.update_parameters()
        self._finalize_scene_init()

    def update_parameters(self):
        pass

    def upload(self, state: np.ndarray = None):
        if not state is None:
            self._buffer.upload(state)

    def render_to(self, target: RenderTarget):
        target.bind()
        self.clear_background(target)
        self._ctx.enable_only(mgl.DEPTH_TEST)
        self._render.render(
            modelview=self.camera.get_view_matrix(),
            projection=self.camera.get_projection_matrix(aspectratio=target.aspectratio),
            camera_position=self.camera.camera_position(),
        )

        self.render_axes(target)

    def list_properties(self):
        return {
            "Charge": self._render.list_properties(),
            **super().list_properties(),
            }


##########################################################################################################


class FlowScene(EmptyScene):
    """Show isosurface of director field."""

    @classmethod
    def supports_system(cls, system: System) -> bool:
        return int(getattr(system, "dim", 0)) == 3

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

        if self._system.dim != 3:
            raise NotImplementedError("Only 3D systems are supported.")

        self._buffer = Field3DBuffer(mglctx=self._ctx, system=self._system)
        self._coordinates_buffer = Field3DBuffer(mglctx=self._ctx, system=self._system)

        idx = self._system.get_inner_simplices()
        self._idx = IntBuffer(mglctx=self._ctx, count=4 * idx.shape[0])
        self._idx.upload(idx.flatten())

        self._render = FlowRender(
            mglctx=self._ctx,
            system=self._system,
            stateBuffer=self._buffer,
            indexBuffer=self._idx,
            coordinatesBuffer=self._coordinates_buffer,
        )

        self.update_parameters()
        self._finalize_scene_init()

    def update_parameters(self):
        pass

    @property
    def surface_render(self):
        return self._render

    def upload(self, state: np.ndarray = None):
        if not state is None:
            self._buffer.upload(state)

    def render_to(self, target: RenderTarget):
        target.bind()
        self.clear_background(target)
        self._ctx.enable_only(mgl.DEPTH_TEST)
        self._render.render(
            modelview=self.camera.get_view_matrix(),
            projection=self.camera.get_projection_matrix(aspectratio=target.aspectratio),
            camera_position=self.camera.camera_position(),
        )

        self.render_axes(target)

    def list_properties(self):
        return {
            "Level surface": self._render.list_properties(),
            **super().list_properties(),
            }


##########################################################################################################


class LevelScene(EmptyScene):
    """Show isusurface of the director field."""

    @classmethod
    def supports_system(cls, system: System) -> bool:
        return int(getattr(system, "dim", 0)) == 3

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

        if self._system.dim != 3:
            raise NotImplementedError("Only 3D systems are supported.")

        self._buffer = Field3DBuffer(mglctx=self._ctx, system=self._system)
        self._coordinates_buffer = Field3DBuffer(mglctx=self._ctx, system=self._system)

        idx = self._system.get_inner_simplices()
        self._idx = IntBuffer(mglctx=self._ctx, count=4 * idx.shape[0])
        self._idx.upload(idx.flatten())

        self._render = LevelRender(
            mglctx=self._ctx,
            system=self._system,
            stateBuffer=self._buffer,
            indexBuffer=self._idx,
            coordinatesBuffer=self._coordinates_buffer,
        )

        self.update_parameters()
        self._finalize_scene_init()

    def update_parameters(self):
        pass

    def upload(self, state: np.ndarray = None):
        if not state is None:
            self._buffer.upload(state)

    def render_to(self, target: RenderTarget):
        target.bind()
        self.clear_background(target)
        self._ctx.enable_only(mgl.DEPTH_TEST)
        self._render.render(
            modelview=self.camera.get_view_matrix(),
            projection=self.camera.get_projection_matrix(aspectratio=target.aspectratio),
            camera_position=self.camera.camera_position(),
        )

        self.render_axes(target)

    def list_properties(self):
        return {
            "Level surface": self._render.list_properties(),
            **super().list_properties(),
            }


##########################################################################################################


class ScalarLevelScene(EmptyScene):
    @classmethod
    def supports_system(cls, system: System) -> bool:
        return int(getattr(system, "dim", 0)) == 3

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

        if self._system.dim!=3:
            raise NotImplementedError("Only 3D systems are supported.")
        
        # self._buffer = Field3DBuffer(mglctx=self._ctx, system=self._system)
        self._coordinates_buffer = Field3DBuffer(mglctx=self._ctx, system=self._system)

        idx = self._system.get_inner_simplices()
        self._idx = IntBuffer(mglctx=self._ctx, count=4*idx.shape[0])
        self._idx.upload(idx.flatten())

        self._render = ScalarLevelRender(mglctx=self._ctx, system=self._system, indexBuffer=self._idx, 
            #stateBuffer=self._buffer,
            coordinatesBuffer=self._coordinates_buffer)

        self.update_parameters()
        self._finalize_scene_init()

    def update_parameters(self):
        pass

    def upload(self, state=None):
        if not state is None:
            self._render.upload(state)

    def render_to(self, target: RenderTarget):
        target.bind()
        self.clear_background(target)
        self._ctx.enable_only(mgl.DEPTH_TEST)
        self._render.render(modelview=self.camera.get_view_matrix()
            , projection=self.camera.get_projection_matrix(aspectratio=target.aspectratio)
            , camera_position=self.camera.camera_position()
        )

        self.render_axes(target)


    def list_properties(self):
        return {
            "Level surface": self._render.list_properties(),
            **super().list_properties(),
        }

##########################################################################################################


class AstraScene(EmptyScene):
    """Show multiple isusurfaces of the director field."""

    @classmethod
    def supports_system(cls, system: System) -> bool:
        return int(getattr(system, "dim", 0)) == 3

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

        if self._system.dim != 3:
            raise NotImplementedError("Only 3D systems are supported.")

        self._buffer = Field3DBuffer(mglctx=self._ctx, system=self._system)
        self._coordinates_buffer = Field3DBuffer(mglctx=self._ctx, system=self._system)

        idx = self._system.get_inner_simplices()
        self._idx = IntBuffer(mglctx=self._ctx, count=4 * idx.shape[0])
        self._idx.upload(idx.flatten())

        self._render = AstraRender(
            mglctx=self._ctx,
            system=self._system,
            stateBuffer=self._buffer,
            indexBuffer=self._idx,
            coordinatesBuffer=self._coordinates_buffer,
        )

        self.update_parameters()
        self._finalize_scene_init()

    def update_parameters(self):
        pass

    def upload(self, state: np.ndarray = None):
        if not state is None:
            self._buffer.upload(state)

    def render_to(self, target: RenderTarget):
        target.bind()
        self.clear_background(target)
        self._ctx.enable_only(mgl.DEPTH_TEST)
        self._render.render(
            modelview=self.camera.get_view_matrix(),
            projection=self.camera.get_projection_matrix(aspectratio=target.aspectratio),
            camera_position=self.camera.camera_position(),
        )

        self.render_axes(target)

    @property
    def surface_render(self):
        return self._render

    def list_properties(self):
        return {
            "Astra": self._render.list_properties(),
            **super().list_properties(),
        }


##########################################################################################################


class IsolinesScene(EmptyScene):
    """Show isolines of the director field."""

    @classmethod
    def supports_system(cls, system: System) -> bool:
        return int(getattr(system, "dim", 0)) == 3

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

        if self._system.dim != 3:
            raise NotImplementedError("Only 3D systems are supported.")

        self._buffer = Field3DBuffer(mglctx=self._ctx, system=self._system)
        self._coordinates_buffer = Field3DBuffer(mglctx=self._ctx, system=self._system)

        idx = self._system.get_inner_simplices()
        self._idx = IntBuffer(mglctx=self._ctx, count=4 * idx.shape[0])
        self._idx.upload(idx.flatten())

        self._render = IsolinesRender(
            mglctx=self._ctx,
            system=self._system,
            stateBuffer=self._buffer,
            indexBuffer=self._idx,
            coordinatesBuffer=self._coordinates_buffer,
        )

        self.update_parameters()
        self._finalize_scene_init()

    @property
    def isolines_render(self):
        return self._render

    def update_parameters(self):
        pass

    def upload(self, state: np.ndarray = None):
        if not state is None:
            self._buffer.upload(state)

    def render_to(self, target: RenderTarget):
        target.bind()
        self.clear_background(target)
        self._ctx.enable_only(mgl.DEPTH_TEST)
        self._ctx.line_width = 3
        self._render.render(
            modelview=self.camera.get_view_matrix(),
            projection=self.camera.get_projection_matrix(aspectratio=target.aspectratio),
            camera_position=self.camera.camera_position(),
        )

        self.render_axes(target)

    @property
    def surface_render(self):
        return self._render

    def list_properties(self):
        return {
            "Isolines": self._render.list_properties(),
            **super().list_properties(),
        }


##########################################################################################################

# List of all available scenes.
SCENES = {
    "Layer": LayerScene,
    "Isolines": IsolinesScene,
    "Level surface": FlowScene,
    "Umbilique surface": ScalarLevelScene,
    "Astra": AstraScene,
    "Layer and isosurface": LayerAndSurfaceScene,
    "Layer and isolines": LayerAndIsolinesScene,
    "Vectors": VectorScene,
    "Cube": CubeScene,
    "Top. charge": ChargeScene,
    "Ball": SphereDotScene,
}


def print_scenes():
    print("Scenes:", ", ".join(SCENES))
